__all__ = ['Trajectory', 'WindowIterator']

from collections import deque
from itertools import islice, chain
from os.path import isfile
from typing import Callable, Optional, Union

import numpy as np
from numpy.typing import NDArray

from dynasor.trajectory.atomic_indices import parse_gromacs_index_file
from dynasor.trajectory.ase_trajectory_reader import ASETrajectoryReader
from dynasor.trajectory.extxyz_trajectory_reader import ExtxyzTrajectoryReader
from dynasor.trajectory.lammps_trajectory_reader import LammpsTrajectoryReader
from dynasor.trajectory.mdanalysis_trajectory_reader import MDAnalysisTrajectoryReader
from dynasor.trajectory.trajectory_frame import TrajectoryFrame
from dynasor.logging_tools import logger


class Trajectory:
    """Instances of this class hold trajectories in a format suitable for
    the computation of correlation functions.  They behave as
    iterators, where each step returns the next frame as a
    :class:`TrajectoryFrame` object.  The latter hold information
    regarding atomic positions, types, and velocities.

    Parameters
    ----------
    filename
        Name of input file.
    trajectory_format
        Type of trajectory. Possible values are:
        ``'lammps_internal'``, ``'extxyz'``, ``'ase'`` or one of the formats supported by
        `MDAnalysis <https://www.mdanalysis.org/>`_ (except for ``'lammpsdump'``,
        which can be called by specifying ``'lammps_mdanalysis'`` to avoid ambiguity)
    atomic_indices
        Specify which indices belong to which atom type. Can be
        (1) a dictionary where the keys specify the species and the values
        are a list of atomic indices,
        (2) ``'read_from_trajectory'``, in which case the species are read from the trajectory or
        (3) the path to a gromacs index file.
    length_unit
        Length unit of trajectory (``'Angstrom'``, ``'nm'``, ``'pm'``, ``'fm'``). Necessary for
        correct conversion to internal dynasor units if the trajectory file does not contain unit
        information.
        If no length unit is specified and the reader cannot read units from the trajectory,
        Angstrom is assumed.
    time_unit
        Time unit of trajectory (``'fs'``, ``'ps'``, ``'ns'``). Necessary for correct conversion to
        internal dynasor units if the trajectory file does not contain unit information.
        If no time unit is specified and the reader cannot read units from the trajectory, fs is
        assumed.
    frame_start
        First frame to read; must be larger or equal ``0``.
    frame_stop
        Last frame to read. By default (``None``) the entire trajectory is read.
    frame_step
        Read every :attr:`frame_step`-th step of the input trajectory.
        By default (``1``) every frame is read. Must be larger than ``0``.

    """
    def __init__(
            self,
            filename: str,
            trajectory_format: str,
            atomic_indices: Optional[Union[str, dict[str, list[int]]]] = None,
            length_unit: Optional[str] = None,
            time_unit: Optional[str] = None,
            frame_start: Optional[int] = 0,
            frame_stop: Optional[int] = None,
            frame_step: Optional[int] = 1
    ):

        if frame_start < 0:
            raise ValueError('frame_start should be positive')
        if frame_step < 0:
            raise ValueError('frame_step should be positive')

        self._frame_start = frame_start
        self._frame_step = frame_step
        self._frame_stop = frame_stop

        # setup trajectory reader
        if not isfile(filename):
            raise IOError(f'File {filename} does not exist')
        self._filename = filename

        if trajectory_format == 'lammps_internal':
            reader = LammpsTrajectoryReader
        elif trajectory_format == 'extxyz':
            reader = ExtxyzTrajectoryReader
        elif trajectory_format == 'lammps_mdanalysis':
            reader = MDAnalysisTrajectoryReader
            trajectory_format = 'lammpsdump'
        elif trajectory_format == 'ase':
            reader = ASETrajectoryReader
        elif trajectory_format == 'lammps':
            raise IOError('Ambiguous trajectory format, '
                          'did you mean lammps_internal or lammps_mdanalysis?')
        else:
            reader = MDAnalysisTrajectoryReader

        logger.debug(f'Using trajectory reader: {reader.__name__}')
        if reader == MDAnalysisTrajectoryReader:
            self._reader_obj = reader(self._filename, trajectory_format,
                                      length_unit=length_unit, time_unit=time_unit)
        else:
            self._reader_obj = reader(self._filename, length_unit=length_unit, time_unit=time_unit)

        # Get two frames to set cell etc.
        frame0 = next(self._reader_obj)
        frame1 = next(self._reader_obj)
        self._cell = frame0.cell
        self._n_atoms = frame0.n_atoms

        # Make sure cell is not changed during consecutive frames
        if not np.allclose(frame0.cell, frame1.cell):
            raise ValueError('The cell changes between the first and second frame. '
                             'The concept of q-points becomes muddy if the simulation cell is '
                             'changing, such as during NPT MD simulations, so trajectories where '
                             'the cell changes are not supported by dynasor.')

        # setup iterator slice (reuse frame0 and frame1 via chain)
        self.number_of_frames_read = 0
        self.current_frame_index = 0
        self._reader_iter = islice(chain([frame0, frame1], self._reader_obj),
                                   self._frame_start, self._frame_stop, self._frame_step)

        # setup atomic indices
        if atomic_indices is None:  # Default behaviour
            atomic_indices = {'X': np.arange(0, self.n_atoms)}
        elif isinstance(atomic_indices, str):  # Str input
            if atomic_indices == 'read_from_trajectory':
                if frame0.atom_types is None:
                    raise ValueError('Could not read atomic indices from the trajectory.')
                else:
                    uniques = np.unique(frame0.atom_types)
                    atomic_indices = {str(uniques[i]):
                                      (frame0.atom_types == uniques[i]).nonzero()[0]
                                      for i in range(len(uniques))}
            else:
                atomic_indices = parse_gromacs_index_file(atomic_indices)
        elif isinstance(atomic_indices, dict):  # dict input
            pass
        else:
            raise ValueError('Could not understand atomic_indices.')
        self._atomic_indices = atomic_indices

        # sanity checks for atomic_indices
        for key, indices in self._atomic_indices.items():
            if np.max(indices) > self.n_atoms:
                raise ValueError('maximum index in atomic_indices exceeds number of atoms')
            if np.min(indices) < 0:
                raise ValueError('minimum index in atomic_indices is negative')
            if '_' in key:
                # Since '_' is what we use to distinguish atom types in the results, e.g. Sqw_Cs_Pb
                raise ValueError('The char "_" is not allowed in atomic_indices.')

        # log info on trajectory and atom types etc
        logger.info(f'Trajectory file: {self.filename}')
        logger.info(f'Total number of particles: {self.n_atoms}')
        logger.info(f'Number of atom types: {len(self.atom_types)}')
        for atom_type, indices in self._atomic_indices.items():
            logger.info(f'Number of atoms of type {atom_type}: {len(indices)}')
        logger.info(f'Simulation cell (in Angstrom):\n{str(self._cell)}')

    def __iter__(self):
        return self

    def __next__(self):
        frame = next(self._reader_iter)
        logger.debug(f'Read frame #{frame.frame_index}')
        new_frame = TrajectoryFrame(
            self.atomic_indices, frame.frame_index, frame.positions, frame.velocities)
        self.number_of_frames_read += 1
        self.current_frame_index = frame.frame_index
        return new_frame

    def __str__(self) -> str:
        s = ['Trajectory']
        s += ['{:12} : {}'.format('filename', self.filename)]
        s += ['{:12} : {}'.format('natoms', self.n_atoms)]
        s += ['{:12} : {}'.format('frame_start', self._frame_start)]
        s += ['{:12} : {}'.format('frame_stop', self._frame_stop)]
        s += ['{:12} : {}'.format('frame_step', self.frame_step)]
        s += ['{:12} : {}'.format('frame_index', self.current_frame_index)]
        s += ['{:12} : [{}\n                {}\n                {}]'
              .format('cell', self.cell[0], self.cell[1], self.cell[2])]
        return '\n'.join(s)

    def __repr__(self) -> str:
        return str(self)

    def _repr_html_(self) -> str:
        s = [f'<h3>{self.__class__.__name__}</h3>']
        s += ['<table border="1" class="dataframe">']
        s += ['<thead><tr><th style="text-align: left;">Field</th><th>Value</th></tr></thead>']
        s += ['<tbody>']
        s += [f'<tr"><td style="text-align: left;">File name</td><td>{self.filename}</td></tr>']
        s += [f'<tr><td style="text-align: left;">Number of atoms</td><td>{self.n_atoms}</td></tr>']
        s += [f'<tr><td style="text-align: left;">Cell metric</td><td>{self.cell}</td></tr>']
        s += [f'<tr><td style="text-align: left;">Frame step</td><td>{self.frame_step}</td></tr>']
        s += [f'<tr><td style="text-align: left;">Atom types</td><td>{self.atom_types}</td></tr>']
        s += ['</tbody>']
        s += ['</table>']
        return '\n'.join(s)

    @property
    def cell(self) -> NDArray[float]:
        """ Simulation cell """
        return self._cell

    @property
    def n_atoms(self) -> int:
        """ Number of atoms """
        return self._n_atoms

    @property
    def filename(self) -> str:
        """ The trajectory filename """
        return self._filename

    @property
    def atomic_indices(self) -> dict[str, list[int]]:
        """ Return copy of index arrays """
        atomic_indices = dict()
        for name, inds in self._atomic_indices.items():
            atomic_indices[name] = inds.copy()
        return atomic_indices

    @property
    def atom_types(self) -> list[str]:
        return sorted(self._atomic_indices.keys())

    @property
    def frame_step(self) -> int:
        """ Frame to access, trajectory will return every :attr:`frame_step`-th snapshot. """
        return self._frame_step


def consume(iterator, n):
    """ Advance the iterator by :attr:`n` steps. If :attr:`n` is ``None``, consume entirely. """
    # From python.org
    if n is None:
        deque(iterator, maxlen=0)
    else:
        next(islice(iterator, n, n), None)


class WindowIterator:
    """Sliding window iterator.

    Returns consecutive windows (a window is represented as a list
    of objects), created from an input iterator.

    Parameters
    ----------
    itraj
        Trajectory object.
    width
        Length of window (``window_size`` + 1).
    window_step
        Distance between the start of two consecutive window frames.
    element_processor
        Enables processing each non-discarded object; useful if ``window_step >
        width`` and ``map_item`` is expensive (as compared to directly passing
        ``map(fun, itraj)`` as ``itraj``); if ``window_step < width``, you could as
        well directly pass ``map(fun, itraj)``.
    """
    def __init__(self,
                 itraj: Trajectory,
                 width: int,
                 window_step: Optional[int] = 1,
                 element_processor: Optional[Callable] = None):

        self._raw_it = itraj
        if element_processor:
            self._it = map(element_processor, self._raw_it)
        else:
            self._it = self._raw_it
        assert window_step >= 1
        assert width >= 1
        self.width = width
        self.window_step = window_step
        self._window = None

    def __iter__(self):
        return self

    def __next__(self):
        """ Returns next element in sequence. """
        if self._window is None:
            self._window = deque(islice(self._it, self.width), self.width)
        else:
            if self.window_step >= self.width:
                self._window.clear()
                consume(self._raw_it, self.window_step - self.width)
            else:
                for _ in range(min((self.window_step, len(self._window)))):
                    self._window.popleft()
            for f in islice(self._it, min((self.window_step, self.width))):
                self._window.append(f)

        if len(self._window) == 0:
            raise StopIteration

        return list(self._window)
