import itertools
from collections import defaultdict
from typing import Any, Dict, Iterable, Iterator, List, Optional  # noqa: F401

# Third-party imports
import torch
import numpy as np

# First-party imports
from .common import DataEntry, Dataset
from pts.feature.transform import Transformation

DataBatch = Dict[str, Any]


class BatchBuffer:
    def __init__(
        self, batch_size: int, device: torch.device, float_type: np.dtype = np.float32
    ) -> None:
        self._buffers: Dict[Any, List[Any]] = defaultdict(list)
        self.batch_size = batch_size
        self._size = 0
        self.device = device
        self.float_type = float_type

    def add(self, d: Dict[str, List[np.ndarray]]):
        if self._buffers:
            assert self._buffers.keys() == d.keys()
        for k, v in d.items():
            self._buffers[k].append(v)
        self._size += 1

    def __len__(self):
        return self._size

    def next_batch(self) -> DataBatch:
        assert self._size > 0
        n = min(self._size, self.batch_size)
        batch = {k: self.stack(v[:n]) for k, v in self._buffers.items()}
        for key in self._buffers.keys():
            self._buffers[key] = self._buffers[key][n:]
        self._size -= n
        return batch

    def stack(self, xs):
        if isinstance(xs[0], np.ndarray):
            data = np.asarray(xs)
            if data.dtype.kind == "f":
                data = data.astype(self.float_type)
            return torch.from_numpy(data).to(device=self.device, non_blocking=True)
        elif isinstance(xs[0], torch.Tensor):
            return torch.stack(*xs)
        else:
            return xs  # stack all other types as list

    def shuffle(self):
        perm = np.random.permutation(self._size)
        for key in self._buffers.keys():
            li = self._buffers[key]
            self._buffers[key] = [li[i] for i in perm]


class DataLoader(Iterable[DataEntry]):
    """
    An abstract Iterable type for iterating and transforming a dataset,
    in batches of a prescribed size.

    Parameters
    ----------
    dataset
        The dataset from which to load data.
    transform
        A transformation to apply to each entry in the dataset.
    batch_size
        The size of the batches to emit.
    device
        device to use to store data on.
    float_type
        Floating point type to use.
    """

    def __init__(
        self,
        dataset: Dataset,
        transform: Transformation,
        batch_size: int,
        device: torch.device,
        float_type: np.dtype = np.float32,
    ) -> None:
        self.dataset = dataset
        self.transform = transform
        self.batch_size = batch_size
        self.device = device
        self.float_type = float_type


class TrainDataLoader(DataLoader):
    """
    An Iterable type for iterating and transforming a dataset, in batches of a
    prescribed size, until a given number of batches is reached.

    The transformation are applied with in training mode, i.e. with the flag
    `is_train = True`.

    Parameters
    ----------
    dataset
        The dataset from which to load data.
    transform
        A transformation to apply to each entry in the dataset.
    batch_size
        The size of the batches to emit.
    device
        device to use to store data on.
    num_batches_per_epoch
        Number of batches to return in one complete iteration over this object.
    float_type
        Floating point type to use.
    """

    def __init__(
        self,
        dataset: Dataset,
        transform: Transformation,
        batch_size: int,
        device: torch.device,
        num_batches_per_epoch: int,
        float_type: np.dtype = np.float32,
        shuffle_for_training: bool = True,
        num_batches_for_shuffling: int = 10,
    ) -> None:
        super().__init__(dataset, transform, batch_size, device, float_type)
        self.num_batches_per_epoch = num_batches_per_epoch
        self.shuffle_for_training = shuffle_for_training
        self._num_buffered_batches = (
            num_batches_for_shuffling if shuffle_for_training else 1
        )
        self._cur_iter: Optional[Iterator] = None
        self._buffer = BatchBuffer(self.batch_size, device, float_type)

    def _emit_batches_while_buffer_larger_than(self, thresh) -> Iterator[DataBatch]:
        if self.shuffle_for_training:
            self._buffer.shuffle()
        while len(self._buffer) > thresh:
            yield self._buffer.next_batch()

    def _iterate_forever(self, collection: Iterable[DataEntry]) -> Iterator[DataEntry]:
        # iterate forever over the collection, the collection must be non empty
        while True:
            try:
                first = next(iter(collection))
            except StopIteration:
                raise Exception("empty dataset")
            else:
                for x in itertools.chain([first], collection):
                    yield x

    def __len__(self) -> int:
        return self.num_batches_per_epoch

    def __iter__(self) -> Iterator[DataBatch]:
        batch_count = 0
        if self._cur_iter is None:
            self._cur_iter = self.transform(
                self._iterate_forever(self.dataset), is_train=True
            )
        assert self._cur_iter is not None
        while True:
            data_entry = next(self._cur_iter)
            self._buffer.add(data_entry)
            if len(self._buffer) >= self._num_buffered_batches * self.batch_size:
                for batch in self._emit_batches_while_buffer_larger_than(
                    self.batch_size - 1
                ):
                    yield batch
                    batch_count += 1
                    if batch_count >= self.num_batches_per_epoch:
                        return


class InferenceDataLoader(DataLoader):
    """
    An Iterable type for iterating and transforming a dataset just once, in
    batches of a prescribed size.

    The transformation are applied with in inference mode, i.e. with the flag
    `is_train = False`.

    Parameters
    ----------
    dataset
        The dataset from which to load data.
    transform
        A transformation to apply to each entry in the dataset.
    batch_size
        The size of the batches to emit.
    device
        device to use to store data on.
    float_type
        Floating point type to use.
    """

    def __iter__(self) -> Iterator[DataBatch]:
        buffer = BatchBuffer(self.batch_size, self.device, self.float_type)
        for data_entry in self.transform(iter(self.dataset), is_train=False):
            buffer.add(data_entry)
            if len(buffer) >= self.batch_size:
                yield buffer.next_batch()
        if len(buffer) > 0:
            yield buffer.next_batch()