r"""Datasets and data loaders."""

import h5py
import numpy as np
import random
import torch

from bisect import bisect
from numpy import ndarray as Array
from pathlib import Path
from torch import Tensor, Size
from torch.distributions import Distribution
from torch.utils.data import DataLoader, Dataset, IterableDataset
from tqdm import tqdm
from typing import *


__all__ = ['JointLoader', 'H5Dataset']


class IterableJointDataset(IterableDataset):
    r"""Creates an iterable dataset of batched pairs :math:`(\theta, x)`."""

    def __init__(
        self,
        prior: Distribution,
        simulator: Callable,
        batch_shape: Size = (),
        numpy: bool = False,
    ):
        super().__init__()

        self.prior = prior
        self.simulator = simulator
        self.batch_shape = batch_shape
        self.numpy = numpy

    def __iter__(self) -> Iterator[Tuple[Tensor, Tensor]]:
        while True:
            theta = self.prior.sample(self.batch_shape)

            if self.numpy:
                x = self.simulator(theta.detach().cpu().numpy().astype(np.float64))
                x = torch.from_numpy(x).to(theta)
            else:
                x = self.simulator(theta)

            yield theta, x


def JointLoader(
    prior: Distribution,
    simulator: Callable,
    batch_size: int = 2**10,  # 1024
    batched: bool = False,
    numpy: bool = False,
    **kwargs,
) -> DataLoader:
    r"""Creates a data loader of batched pairs :math:`(\theta, x)` generated by
    a prior distribution :math:`p(\theta)` and a simulator.

    The simlator is a stochastic function taking a set of parameters :math:`\theta`
    as input and returning an observation :math:`x` as output, which (implicitely)
    defines a likelihood function :math:`p(x | \theta)`. Together with the prior,
    they form a joint distribution :math:`p(\theta, x) = p(\theta) p(x | \theta)`
    from which the pairs :math:`(\theta, x)` are independently drawn.

    Arguments:
        prior: A prior distribution :math:`p(\theta)`.
        simulator: A callable simulator.
        batch_size: The batch size of the generated pairs.
        batched: Whether the simulator accepts batched inputs or not.
        numpy: Whether the simulator requires NumPy or PyTorch inputs.
        kwargs: Keyword arguments passed to :class:`torch.utils.data.DataLoader`.

    Returns:
        An infinite data loader of batched pairs :math:`(\theta, x)`.

    Example:
        >>> loader = joint_loader(prior, simulator, numpy=True, num_workers=4)
        >>> for theta, x in loader:
        ...     theta, x = theta.cuda(), x.cuda()
        ...     something(theta, x)
    """

    return DataLoader(
        IterableJointDataset(
            prior,
            simulator,
            batch_shape=(batch_size,) if batched else (),
            numpy=numpy,
        ),
        batch_size=None if batched else batch_size,
        **kwargs,
    )


class H5Dataset(object):
    r"""Creates a dataset of pairs :math:`(\theta, x)` from HDF5 files.

    As a :class:`torch.utils.data.Dataset`, :class:`H5Dataset` implements the methods
    :meth:`__len__` and :meth:`__getitem__`. However, as it can be slow to load pairs
    from disk one by one when iterating over the dataset, it also implements a custom
    :meth:`__iter__` method. This method loads several contiguous chunks of pairs at
    once, concatenates them, shuffles the result and, finally, splits it into batches.
    This "weak shuffling" procedure greatly improves loading performances, but the
    resulting batch elements are not perfectly independent from each others.

    Important:
        To take advantage of the custom :meth:`__iter__` method, :class:`H5Dataset`
        instances should not be wrapped inside a :class:`torch.utils.data.DataLoader`
        when iterating over the dataset.

    Arguments:
        files: HDF5 files containing pairs :math:`(\theta, x)`.
        batch_size: The size of the batches.
        chunk_size: The size of the contiguous chunks.
        group_size: The number of chunks loaded at once.
        pin_memory: Whether the batches reside in CUDA pinned memory or not.
        shuffle: Whether the pairs are shuffled when iterating.
        seed: A seed to initialize the internal RNG used for shuffling.

    Example:
        >>> dataset = H5Dataset('data.h5', batch_size=256, shuffle=True)
        >>> theta, x = dataset[0]
        >>> theta
        tensor([-0.1215, -1.3641,  0.7233, -1.2150, -1.9263])
        >>> for theta, x in dataset:
        ...     theta, x = theta.cuda(), x.cuda()
        ...     something(theta, x)
    """

    def __init__(
        self,
        *files: Union[str, Path],
        batch_size: int = 2**10,  # 1024
        chunk_size: int = 2**12,  # 4096
        group_size: str = 2**4,  # 16
        pin_memory: bool = False,
        shuffle: bool = True,
        seed: int = None,
    ):
        super().__init__()

        self.files = list(map(h5py.File, files))
        self.stops = np.cumsum([len(f['x']) for f in self.files])

        self.batch_size = batch_size
        self.chunk_size = chunk_size
        self.group_size = group_size

        self.pin_memory = pin_memory

        self.shuffle = shuffle
        self.rng = np.random.default_rng(seed)

    def __del__(self) -> None:
        for f in self.files:
            f.close()

    def __len__(self) -> int:
        return self.stops[-1]

    def __getitem__(self, i: int) -> Tuple[Tensor, Tensor]:
        i = i % len(self)
        j = bisect(self.stops, i)
        if j > 0:
            i = i - self.stops[j - 1]

        f = self.files[j]

        if 'theta' in f:
            theta = torch.from_numpy(f['theta'][i])
        else:
            theta = None

        x = torch.from_numpy(f['x'][i])

        return theta, x

    def __iter__(self) -> Iterator[Tuple[Tensor, Tensor]]:
        chunks = [
            (i, j, j + self.chunk_size)
            for i, f in enumerate(self.files)
            for j in range(0, len(f['x']), self.chunk_size)
        ]

        if self.shuffle:
            self.rng.shuffle(chunks)

        for l in range(0, len(chunks), self.group_size):
            slices = sorted(chunks[l:l+self.group_size])

            # Load
            theta = np.concatenate([self.files[i]['theta'][j:k] for i, j, k in slices])
            x = np.concatenate([self.files[i]['x'][j:k] for i, j, k in slices])

            # Shuffle
            if self.shuffle:
                order = self.rng.permutation(len(theta))
                theta, x = theta[order], x[order]

            # Tensor
            theta, x = torch.from_numpy(theta), torch.from_numpy(x)

            if self.pin_memory:
                theta, x = theta.pin_memory(), x.pin_memory()

            # Batches
            yield from zip(
                theta.split(self.batch_size),
                x.split(self.batch_size),
            )

    @staticmethod
    def store(
        pairs: Iterable[Tuple[Array, Array]],
        file: Union[str, Path],
        size: int,
        dtype: np.dtype = np.float32,
        **meta,
    ) -> None:
        r"""Creates an HDF5 file containing pairs :math:`(\theta, x)`.

        The sets of parameters :math:`\theta` are stored in a collection named
        :py:`'theta'` and the observations in a collection named :py:`'x'`.

        Arguments:
            pairs: An iterable over batched pairs :math:`(\theta, x)`.
            file: An HDF5 filename to store pairs in.
            size: The number of pairs to store.
            dtype: The data type to store pairs in.
            meta: Metadata to store in the file.

        Example:
            >>> loader = JointLoader(prior, simulator, batch_size=16)
            >>> H5Dataset.store(loader, 'sim.h5', 4096)
            100%|██████████| 4096/4096 [01:35<00:00, 42.69sample/s]
        """

        # File
        file = Path(file)
        file.parent.mkdir(parents=True, exist_ok=True)

        with h5py.File(file, 'w-') as f:
            ## Attributes
            f.attrs.update(meta)

            ## Datasets
            theta, x = map(np.asarray, next(iter(pairs)))
            theta, x = theta[0], x[0]

            f.create_dataset('theta', (size,) + theta.shape, dtype=dtype)
            f.create_dataset('x', (size,) + x.shape, dtype=dtype)

            ## Store
            with tqdm(total=size, unit='sample') as tq:
                i = 0

                for theta, x in pairs:
                    j = min(i + theta.shape[0], size)

                    f['theta'][i:j] = np.asarray(theta)[:j-i]
                    f['x'][i:j] = np.asarray(x)[:j-i]

                    tq.update(j - i)

                    if j < size:
                        i = j
                    else:
                        break
