from torch.utils.data import Sampler, Dataset, DataLoader, BatchSampler
from typing import List, Optional, Tuple
from pathlib import Path
from torch import Tensor
from tqdm import tqdm
import numpy as np
import itertools
import torch
import h5py


from torch_utilities.common import to_numpy, split_complex

# export list
__all__ = [
    "HDF5Dataset",
    "get_hdf5_dataloader",
    "HDF5OnlineDataset",
    "get_dataset_statistics",
    "plot_dataset_statistics",
]


class WeakShufflingSampler(Sampler):
    def __init__(self, dataset: Dataset, batch_size: int):
        """
        Sampler that implements weak-shuffling.

        E.g. if dataset is [1,2,3,4] with batch_size=2,
             then first batch, [[1,2], [3,4]] then
             shuffle batches -> [[3,4],[1,2]]

        Referenece:
        https://towardsdatascience.com/reading-h5-files-faster-with-pytorch-datasets-3ff86938cc

        Parameters
        ----------
        dataset : Dataset
            Source dataset
        batch_size : int
            Size of the batches
        """
        self.batch_size = batch_size
        self.dataset_length = len(dataset)
        self.n_batches = int(np.ceil((self.dataset_length / self.batch_size)))
        self.batch_ids = torch.randperm(self.n_batches)

    def __len__(self):
        return self.batch_size

    def __iter__(self):
        datalen = self.dataset_length
        for id in self.batch_ids:
            start = id * self.batch_size
            end = (id + 1) * self.batch_size
            end = end if end < datalen else datalen
            selection = torch.arange(start, end)
            for index in selection:
                yield int(index)


class HDF5Dataset(Dataset):
    def __init__(
        self,
        dataset_path: Path,
        data_layout: List[str],
    ) -> None:
        """
        Dataset supporting HDF5 format.

        The following constraints should be satisfied
        by the HDF5 file for a correct behaviour:
        - Multiple single-level groups (e.g. /group1, /group2, ...)
        - The batch size (first dimension) of each dataset is the same
        - Each group has len(data_layout) dataset inside

        Parameters
        ----------
        dataset_path : Path
            Path to the .hdf5 file
        data_layout : List[str]
            List describing the layout of the data
            inside each group. For an input-label1-label2
            dataset the list would be ["input", "label1", "label2"]
        """
        super().__init__()
        self.dataset_path = dataset_path.absolute()
        self.data_layout = data_layout
        self.dataset_file = h5py.File(self.dataset_path, "r")
        self.groups = list(self.dataset_file.keys())
        self.group_batch_len = self.dataset_file[self.groups[0]][
            self.data_layout[0]
        ].shape[0]
        self.data_layout = data_layout
        self._cache = None
        self._cache_idx = None

    def __len__(self):
        return len(self.groups) * self.group_batch_len

    def __getitem__(self, idx) -> List[Tensor]:
        # error handling
        err_msg = None
        if not isinstance(idx, list):
            err_msg = "HDF5Dataset must be sampled by a BatchLoader"
        elif len(idx) > self.group_batch_len:
            err_msg = (
                "Reduce the batch size to be less than the batch size of the groups"
            )
        elif self.group_batch_len % len(idx) != 0:
            err_msg = (
                "Modify the batch size to be a divider of the HDF5 group batch size"
            )
        if err_msg is not None:
            raise RuntimeError(err_msg)

        # cache update
        if not self._in_cache(idx[0]):
            self._update_cache(idx[0])

        # using the cache
        gbl = self.group_batch_len
        a, b = idx[0] % gbl, (idx[-1] % gbl) + 1
        data = [x[a:b] for x in self._cache]

        return data

    def _update_cache(self, idx: int) -> None:
        """
        Caches a group in memory.

        Parameters
        ----------
        idx : int
            Index of an element of the group
        """
        # getting the starting index of a group
        self._cache_idx = idx

        # updating the cache
        del self._cache
        g_idx = idx // self.group_batch_len
        g = self.dataset_file[self.groups[g_idx]]
        cast = lambda x: torch.from_numpy(np.array(x)).to("cpu")
        data = [cast(g[k]) for k in self.data_layout]
        self._cache = data

    def _in_cache(self, idx: int) -> bool:
        """
        Checks if an index is inside the cache.

        Parameters
        ----------
        idx : int
            Target index

        Returns
        -------
        bool
            True if the element is inside the cache
        """
        c_idx = self._cache_idx
        flag = (c_idx is not None) and (
            idx >= c_idx and idx < (c_idx + self.group_batch_len)
        )
        return flag


def get_hdf5_dataloader(
    dataset_path: Path,
    data_layout: List[str],
    batch_size: int = None,
    **dataloader_kwargs: dict,
) -> DataLoader:
    """
    Create a dataloader binded to a HDF5Dataset.

    Parameters
    ----------
    dataset : Path
        Path to the HDF5 Dataset
    data_layout : List[str]
            List describing the layout of the data
            inside each group. For an input-label1-label2
            dataset the list would be ["input", "label1", "label2"]
    batch_size : int, optional
        Batch size of the dataloader, by default HDF5Dataset.group_batch_len
    dataloader_kwargs : dict, optional
        DataLoader arguments, by default {}
    """
    if batch_size is None:
        batch_size = dataset.group_batch_len

    if dataloader_kwargs is None:
        dataloader_kwargs = {}

    dataset = HDF5Dataset(dataset_path, data_layout)

    sampler = BatchSampler(
        sampler=WeakShufflingSampler(
            dataset=dataset,
            batch_size=batch_size,
        ),
        batch_size=batch_size,
        drop_last=False,
    )

    dataloader = DataLoader(
        dataset=dataset,
        sampler=sampler,
        **dataloader_kwargs,
    )

    return dataloader


class HDF5OnlineDataset(Dataset):
    def __init__(
        self,
        dataset_paths: List[Path],
        data_layouts: List[List[str]],
        batch_size: int = 16,
        total_items: int = 1024,
        overfit_mode: bool = False,
    ) -> None:
        """
        Dataset based on online generation.
        Data can be combined after retrieving it
        from source HDF5Datasets.

        Parameters
        ----------
        dataset_paths : List[Path]
            List of paths to the HDF5Datasets h5
        data_layouts : List[List[str]]
            Data layouts for all the HDF5Datasets
        batch_size : int, optional
            Batch size of each raw data, by default 16
        total_items : int, optional
            Total item in an epoch, by default 1024
        overfit_mode : bool, optional
            If True, always returns the first batch at every iteration
        """
        super(HDF5OnlineDataset).__init__()
        self.datasets_paths = dataset_paths
        self.data_layouts = data_layouts
        self.batch_size = batch_size
        self.total_items = total_items
        self.overfit_mode = overfit_mode

        # error handling
        err_msg = f"total_items ({total_items}) must be divisible by batch_size ({batch_size})"
        assert (self.total_items % self.batch_size) == 0, err_msg

        # static datasets
        self.source_datasets = self._get_datasets()

        # handling overfit mode
        self.overfit_cache = None

    def __len__(self) -> int:
        return self.total_items // self.batch_size

    def transform(self, raw_data: List[Tensor]) -> List[Tensor]:
        """
        Override this method to process the raw data.

        Parameters
        ----------
        raw_data : List[Tensor]
            Raw data from the HDF5Datasets

        Returns
        -------
        List[Tensor]
            Processed data
        """
        return raw_data

    def _get_rand_batch(self, dataset: HDF5Dataset) -> List[Tensor]:
        """
        Get a random batch from an HDF5Dataset

        Parameters
        ----------
        dataset : HDF5Dataset
            Target HDF5Dataset

        Returns
        -------
        List[Tensor]
            Batches from the HDF5Datasets
        """
        # getting start and end points
        ds_len = len(dataset)
        total_batches = ds_len // self.batch_size
        r_int = np.random.randint(0, total_batches)
        start = r_int * self.batch_size
        end = (r_int + 1) * self.batch_size
        indices = list(range(start, end))

        # getting the random selection
        ds_outs = dataset[indices]
        return ds_outs

    def __getitem__(self, _) -> List[Tensor]:
        if self.overfit_cache is not None:
            # overfit cache enabled
            return self.overfit_cache

        outs = [self._get_rand_batch(ds) for ds in self.source_datasets]
        outs = itertools.chain.from_iterable(outs)
        outs = list(outs)
        outs = self.transform(outs)

        if self.overfit_mode:
            # caching for overfit mode
            self._store_in_cache(outs)

        return outs

    def _get_datasets(self) -> List[HDF5Dataset]:
        """
        Gets the static datasets.

        Returns
        -------
        List[HDF5Dataset]
            List of static datasets
        """
        ds = [HDF5Dataset(p, l) for p, l in zip(self.datasets_paths, self.data_layouts)]
        return ds

    def _store_in_cache(self, data: List[Tensor]) -> None:
        """
        Stores the data in the overfit cache.

        Parameters
        ----------
        data : List[Tensor]
            Data to store in the cache
        """
        self.overfit_cache = data


def get_dataset_statistics(
    dataloader: DataLoader,
    iterations: int,
    n_hist_bins: int = 100,
    verbose: bool = False,
) -> List[Tuple[np.ndarray, np.ndarray, float, float, float]]:
    """
    Calculate the statistics for the data obtained
    from a dataset.

    Parameters
    ----------
    dataloader : DataLoader
        DataLoader to iterate
    iterations : int
        Number of iterations
    n_hist_bins : int, optional
        Number of histogram bins
    verbose : bool, optional
        If True shows a progress bar for the iterations,
        by default False

    Returns
    -------
    List[Tuple[np.ndarray, np.ndarray, float, float, float, float]]
        List of statistics for each dataset output
        each tuple contains
        (hist_vals, hist_bins, min, max, mean, var)
    """
    ds = iter(dataloader)
    _range = tqdm(range(iterations)) if verbose else range(iterations)
    outs = [ds.__next__() for i in _range]
    n_items = len(outs[0])

    stats = []
    for i in range(n_items):
        x = torch.stack([o[i] for o in outs])
        x = x.flatten()
        x = x.float()
        x = split_complex(x) if torch.is_complex(x) else x
        hist_vals, hist_bins = [
            to_numpy(y) for y in x.histogram(n_hist_bins, density=True)
        ]
        hist_vals /= n_hist_bins
        _min = x.min().item()
        _max = x.max().item()
        _mean = x.mean().item()
        _var = x.var().item()
        stats.append((hist_vals, hist_bins, _min, _max, _mean, _var))

    return stats


def plot_dataset_statistics(
    dataloader: DataLoader,
    iterations: int,
    n_hist_bins: int = 100,
    labels: Optional[List[str]] = None,
    save_path: Optional[Path] = None,
    figsize: Tuple[int, int] = (8, 6),
    verbose: bool = False,
) -> List[Tuple[np.ndarray, np.ndarray, float, float, float, float]]:
    """
    Plots the statistics for the data obtained
    from a dataset.

    Parameters
    ----------
    dataloader : DataLoader
        DataLoader to iterate
    iterations : int
        Number of iterations
    n_hist_bins : int, optional
        Number of histogram bins
    labels : Optional[List[str]]
        Names of the dataset outputs
    save_path : Optional[Path]
        If specified saves the plots in the specified folder
    figsize : Tuple[int, int], optional
        Size of the plot, by default (8, 6)
    verbose : bool, optional
        If True shows a progress bar for the iterations,
        by default False


    Returns
    -------
    List[Tuple[np.ndarray, np.ndarray, float, float, float]]
        List of statistics for each dataset output
        each tuple contains
        (hist_vals, hist_bins, min, max, mean, var)
    """
    import matplotlib.pyplot as plt

    all_stats = get_dataset_statistics(
        dataloader,
        iterations,
        n_hist_bins,
        verbose,
    )

    # default labels
    if labels is None:
        labels = [f"feature_{i}" for i in range(len(all_stats))]

    # optional save path
    if save_path is not None:
        save_path = Path(save_path)

    for stats, name in zip(all_stats, labels):
        hist_vals, hist_bins, min, max, mean, var = stats
        bar_width = np.abs(hist_bins[1] - hist_bins[0])
        plt.figure(figsize=figsize)
        plt.title(f"{name} distribution")
        plt.bar(hist_bins[:-1], hist_vals, width=bar_width)
        plt.vlines(x=mean, ymin=0, ymax=1, colors="red", label="mean")
        plt.ylim([0, hist_vals.max()])
        plt.legend()
        plt.grid()
        if save_path is not None:
            plt.savefig(save_path / f"{name}.png")
        else:
            plt.show()
        plt.close()

    return all_stats
