import bisect
import pickle
from abc import ABCMeta, abstractmethod
from collections import deque
from functools import lru_cache
from itertools import accumulate, chain, islice, tee
from pathlib import Path
from typing import Any, Callable, Iterable, Iterator, List, Tuple, Union

import arrayfiles
from _collections_abc import Sequence, _check_methods


class DatasetMixin(metaclass=ABCMeta):

    __slots__ = ()

    def __getitem__(self, index: Union[int, slice]) -> Union[Any, List[Any]]:
        if isinstance(index, slice):
            start, stop, step = index.indices(len(self))
            return [self.get_example(i) for i in range(start, stop, step)]

        if index >= 0:
            if index >= len(self):
                raise IndexError(f'{self.__class__.__name__} object index out of range')
        else:
            if index < - len(self):
                raise IndexError(f'{self.__class__.__name__} object index out of range')
            index += len(self)

        return self.get_example(index)

    @abstractmethod
    def __iter__(self) -> Iterator[Any]:
        while False:
            yield None

    @abstractmethod
    def get_example(self, i: int) -> Any:
        raise IndexError

    @abstractmethod
    def __len__(self) -> int:
        return 0

    @classmethod
    def __subclasshook__(cls, C):
        if cls is DatasetMixin:
            return _check_methods(C, '__iter__', 'get_example', '__len__')
        return NotImplemented


DatasetMixin.register(Sequence)
DatasetMixin.register(arrayfiles.TextFile)


class Dataset(DatasetMixin):
    """Dataset wrapping ``DatasetMixin`` object.

    Args:
        dataset (DatasetMixin): ``Sequence``, ``arrayfiles.TextFile``, or ``arrayfiles.CsvFile`` object.
    """

    def __init__(self,
                 dataset: DatasetMixin) -> None:
        assert isinstance(dataset, DatasetMixin)

        self._dataset = dataset
        self._length = None

    def __iter__(self) -> Iterator[Any]:
        yield from self._dataset

    def get_example(self, i: int) -> Any:
        return self._dataset[i]

    def __len__(self) -> int:
        if self._length is None:
            self._length = len(self._dataset)
        return self._length

    def __add__(self, other: 'Dataset') -> 'ConcatDataset':
        return ConcatDataset(self, other)

    def map(self, map_func: Callable[[Any], Any]) -> 'MapDataset':
        """Applies a function across the examples of this dataset.

        Args:
            map_func (Callable[[Any], Any]): A function to apply.

        Returns ('MapDataset'):
            The dataset applied the function.
        """
        return MapDataset(self, map_func)

    def flat_map(self, map_func: Callable[[Any], Any]) -> 'IterableDataset':
        """Applies a function across the examples of this dataset and then flattens the result.

        Args:
            map_func (Callable[[Any], Any]): A function to apply.

        Returns ('IterableDataset'):
            The dataset applied the function and flattened.
        """
        return IterableDataset(lineflow_flat_map(map_func, self, lazy=True))

    def filter(self, predicate: Callable[[Any], bool]) -> 'IterableDataset':
        """Filters this dataset by a predicate function.

        Args:
            predicate (Callable[[Any], bool]): A predicate function.

        Returns ('IterableDataset'):
            The dataset containing the examples for which ``predicate`` returns ``True``.
        """
        return IterableDataset(lineflow_filter(predicate, self, lazy=True))

    def window(self, window_size: int, shift: int = None) -> 'IterableDataset':
        """Combines input examples into a dataset of windows.

        Args:
            window_size (int): the number of examples of the input dataset to combine into a window.
            shift (int, optional): The forward shift of the sliding window in each iteration.

        Returns ('IterableDataset'):
            The dataset of windows.
        """
        return IterableDataset(lineflow_window(self, window_size, shift, lazy=True))

    def all(self) -> List[Any]:
        """Takes all examples from the dataset.

        Returns (List[Any]):
            The list of the examples in the dataset.
        """
        return list(self)

    def take(self, n: int) -> List[Any]:
        """Takes the first n examples from the dataset.

        Args:
            n (int): the number of examples to take.

        Returns (List[Any]):
            The list of the ``n`` examples.
        """
        return list(islice(self, n))

    def first(self) -> Any:
        """Takes the first example from the dataset.

        Returns (Any):
            The first example in the dataset.
        """
        return next(iter(self))

    def save(self, filename: str) -> 'CacheDataset':
        """Evaluates the datasets and save it as pickle.

        Args:
            filename (str): The name of the pickle file.

        Returns ('CacheDataset'):
            The evaluated dataset.
        """
        path = Path(filename)
        if path.exists():
            print(f'Loading data from {filename}...')
            with path.open('rb') as f:
                cache = pickle.load(f)
        else:
            if not path.parent.exists():
                path.parent.mkdir(parents=True)
            print(f'Saving data to {filename}...')
            cache = list(self)
            with path.open('wb') as f:
                pickle.dump(cache, f)
        return CacheDataset(cache)


class IterableDataset(Dataset):
    def __init__(self, iterable: Iterable) -> None:
        self._length = None
        self._iterable = iterable
        self._computed = False

    @lru_cache()
    def _get_dataset(self) -> List[Any]:
        self._computed = True
        return list(self._iterable)

    @property
    def _dataset(self) -> List[Any]:
        return self._get_dataset()

    def __iter__(self) -> Iterator[Any]:
        if self._computed:
            yield from self._dataset
        else:
            iterable, self._iterable = tee(self._iterable)
            yield from iterable

    def get_example(self, i: int) -> Any:
        return super(IterableDataset, self).get_example(i)

    def __len__(self) -> int:
        return super(IterableDataset, self).__len__()


class ConcatDataset(Dataset):
    def __init__(self, *datasets: List[DatasetMixin]) -> None:
        assert all(isinstance(d, DatasetMixin) for d in datasets)

        self._datasets = datasets
        self._length = None

    @lru_cache()
    def _get_lengths(self) -> List[int]:
        return list(accumulate(len(d) for d in self._datasets))

    @property
    def _lengths(self) -> List[int]:
        return self._get_lengths()

    @lru_cache()
    def _get_offsets(self) -> List[int]:
        offsets = [0]
        offsets.extend(self._lengths[:-1])
        return offsets

    @property
    def _offsets(self) -> List[int]:
        return self._get_offsets()

    def __iter__(self) -> Iterator[Any]:
        for d in self._datasets:
            yield from d

    def get_example(self, i: int) -> Any:
        j = bisect.bisect_right(self._lengths, i)
        return self._datasets[j][i - self._offsets[j]]

    def __len__(self) -> int:
        if self._length is None:
            self._length = self._lengths[-1]
        return self._length


class ZipDataset(Dataset):
    def __init__(self, *datasets: List[DatasetMixin]) -> None:
        assert all(isinstance(d, DatasetMixin) for d in datasets)
        self._datasets = datasets
        self._length = None

    def __iter__(self) -> Iterator[Tuple[Any]]:
        yield from zip(*self._datasets)

    def get_example(self, i: int) -> Tuple[Any]:
        return tuple(d[i] for d in self._datasets)

    def __len__(self) -> int:
        if self._length is None:
            self._length = min(len(d) for d in self._datasets)
        return self._length


class MapDataset(Dataset):
    def __init__(self,
                 dataset: DatasetMixin,
                 map_func: Callable[[Any], Any]) -> None:
        assert callable(map_func)

        self._map_func = map_func

        super(MapDataset, self).__init__(dataset)

    def __iter__(self) -> Iterator[Any]:
        yield from map(self._map_func, self._dataset)

    def get_example(self, i: int) -> Any:
        return self._map_func(self._dataset[i])


class CacheDataset(Dataset):
    def __init__(self, cache: List[Any]) -> None:
        super(CacheDataset, self).__init__(cache)

        self._length = len(cache)


def lineflow_concat(*datasets: List[DatasetMixin]) -> ConcatDataset:
    return ConcatDataset(*datasets)


def lineflow_zip(*datasets: List[DatasetMixin]) -> ZipDataset:
    return ZipDataset(*datasets)


def lineflow_filter(
        predicate: Callable[[Any], bool],
        dataset: DatasetMixin,
        lazy: bool = False) -> Union[Iterator[Any], List[Any]]:
    iterator = filter(predicate, dataset)
    if lazy:
        return iterator
    else:
        return list(iterator)


def lineflow_flat_map(
        map_func: Callable[[Iterable[Any]], Any],
        dataset: DatasetMixin,
        lazy: bool = False) -> Union[Iterator[Any], List[Any]]:
    iterator = chain.from_iterable(map(map_func, dataset))
    if lazy:
        return iterator
    else:
        return list(iterator)


def lineflow_window(
        dataset: DatasetMixin,
        window_size: int,
        shift: int = None,
        lazy: bool = False) -> Union[Iterator[Any], List[Any]]:
    shift = shift or window_size

    def generator(dataset, window_size, shift):
        iterator = iter(dataset)
        window = deque([], window_size)
        append = window.append

        for _, x in zip(range(window_size), iterator):
            append(x)
        yield tuple(window)

        i = 0
        for x in iterator:
            append(x)
            i = (i + 1) % shift
            if i % shift == 0:
                yield tuple(window)

        if (i % shift) and (shift - i < window_size):
            popleft = window.popleft
            for _ in range(shift - i):
                popleft()
            yield tuple(window)

    iterator = generator(dataset, window_size, shift)
    if lazy:
        return iterator
    else:
        return list(iterator)


def lineflow_load(filename: str) -> Dataset:
    print(f'Loading data from {filename}...')
    with open(filename, 'rb') as f:
        dataset = pickle.load(f)
    return Dataset(dataset)
