import concurrent.futures
import numpy as np
import io
from itertools import count
from typing import Optional
from dynasor.trajectory.abstract_trajectory_reader import AbstractTrajectoryReader
from dynasor.trajectory.trajectory_frame import ReaderFrame
from ase.io import read
from ase import Atoms


def frame_text_to_atoms(frame_text):
    return read(io.StringIO(frame_text), format='extxyz')


def iter_frame_texts(f):
    while True:
        natoms_line = f.readline()
        if not natoms_line:
            break

        comment_line = f.readline()
        if not comment_line:
            raise ValueError('Unexpected EOF while reading extxyz comment line')

        try:
            natoms = int(natoms_line.strip())
        except ValueError:
            raise ValueError(f'Invalid extxyz atom count line: {natoms_line!r}')

        atom_lines = []
        for _ in range(natoms):
            line = f.readline()
            if not line:
                raise ValueError('Unexpected EOF while reading extxyz atom lines')
            atom_lines.append(line)

        yield natoms_line + comment_line + ''.join(atom_lines)


def iread(f, max_workers: Optional[int] = None) -> Atoms:
    frame_iterator = iter_frame_texts(f)
    with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as ex:
        buff = []
        n_submit = ex._max_workers

        for _ in range(n_submit):
            try:
                frame_text = next(frame_iterator)
                buff.append(ex.submit(frame_text_to_atoms, frame_text))
            except StopIteration:
                break

        while buff:
            res = buff.pop(0)

            try:
                frame_text = next(frame_iterator)
                buff.append(ex.submit(frame_text_to_atoms, frame_text))
            except StopIteration:
                pass

            yield res.result()


class ExtxyzTrajectoryReader(AbstractTrajectoryReader):
    """Read extend xyz trajectory file, typically produced by GPUMD

    This is a naive (and comparatively slow) parallel implementation which
    relies on the ASE xyz reader.

    Parameters
    ----------
    filename
        Name of input file.
    length_unit
        Unit of length for the input trajectory (``'Angstrom'``, ``'nm'``, ``'pm'``, ``'fm'``).
    time_unit
        Unit of time for the input trajectory (``'fs'``, ``'ps'``, ``'ns'``).
    max_workers
        Number of working processes; defaults to ``None``, which means that the number of
        processors on the machine is used.
    """

    def __init__(self,
                 filename: str,
                 length_unit: Optional[str] = None,
                 time_unit: Optional[str] = None,
                 max_workers: Optional[int] = None):

        # setup generator object
        self._fobj = open(filename, 'r')
        self._generator_xyz = iread(self._fobj, max_workers=max_workers)
        self._open = True
        self._frame_index = count(0)

        # set up units
        self.set_unit_scaling_factors(length_unit, time_unit)

    def _get_next(self):
        try:
            atoms = next(self._generator_xyz)
        except StopIteration:
            self._fobj.close()
            self._open = False
            raise
        except Exception:
            self._fobj.close()
            self._open = False
            raise

        self._atom_types = np.array(list(atoms.symbols))
        self._n_atoms = len(atoms)
        self._cell = atoms.cell[:]
        self._x = atoms.positions
        if 'vel' in atoms.arrays:
            self._v = atoms.arrays['vel']
        else:
            self._v = None

    def __iter__(self):
        return self

    def close(self):
        if not self._fobj.closed:
            self._fobj.close()
            self._open = False

    def __next__(self):
        if not self._open:
            raise StopIteration

        self._get_next()

        if self._v is not None:
            frame = ReaderFrame(frame_index=next(self._frame_index),
                                n_atoms=int(self._n_atoms),
                                cell=self.x_factor * self._cell.copy('F'),
                                positions=self.x_factor * self._x,
                                velocities=self.v_factor * self._v,
                                atom_types=self._atom_types
                                )
        else:
            frame = ReaderFrame(frame_index=next(self._frame_index),
                                n_atoms=int(self._n_atoms),
                                cell=self.x_factor * self._cell.copy('F'),
                                positions=self.x_factor * self._x,
                                atom_types=self._atom_types
                                )
        return frame
