# -*- coding:utf-8 -*-

# Copyright xmuspeech (Author: Leo 2021-09-20)
# Copyright TalentedSoft (Author: <NAME> 2021-10-13)

import numpy as np
import random
import sklearn
import pandas as pd
import torch
import torchaudio
from torch.utils.data import RandomSampler
from torch.utils.data import Dataset
from torch.utils.data import DataLoader

torchaudio.set_audio_backend("sox_io")

class CSVAugmentDataset(Dataset):
    def __init__(self, csv_file, sorting="original", max_len=None):
        head = pd.read_csv(csv_file, sep=" ", nrows=0).columns
        assert "ID" in head
        assert "wav" in head
        assert "duration" in head
        assert sorting in ["original", "decending", "ascending","random"]

        data = pd.read_csv(csv_file, sep=" ",header=0)

        if sorting == "decending":
            data = data.sort_values(by=['duration'], ascending=False)
        elif sorting == "ascending":
            data = data.sort_values(by=['duration'], ascending=True)
        elif sorting == "random":
            data = sklearn.utils.shuffle(data)
        else:
            pass

        self.path = data['wav'].values.astype(np.string_)
        if max_len:
            assert max_len > 0.0
            self.lens=data['tot_frame'].values
            self.sr=data['sr'].values

        self.max_len = max_len
        del data


    def __getitem__(self, index):
        wav=str(self.path[index], encoding='utf-8')

        if self.max_len:
            audio_len = self.lens[index]
            max_frame = min(int(self.sr[index] * self.max_len),audio_len)

            start = max(int(random.random()*(audio_len - max_frame)),0)

            wavforms, fs = torchaudio.backend.sox_io_backend.load(wav,num_frames=max_frame, frame_offset=start)

        else:
            wavforms, fs = torchaudio.load(wav)

        return wavforms[0]

    def __len__(self):
        return len(self.path)


def make_dataloader(dataset,  **loader_kwargs):
    """Makes a basic DataLoader.

    Arguments
    ---------
    dataset : Dataset
        The dataset to make a DataLoader.
    **loader_kwargs : dict
        Keyword args to DataLoader, see PyTorch DataLoader for
        options.

    Returns
    -------
    DataLoader
    """

    # Reproducible random sampling
    if loader_kwargs.get("shuffle", False):
        if loader_kwargs.get("sampler") is not None:
            raise ValueError(
                "Cannot specify both shuffle=True and a "
                "sampler in loader_kwargs"
            )
        sampler = ReproducibleRandomSampler(dataset)
        loader_kwargs["sampler"] = sampler
        # Should delete shuffle because you can't set both Sampler and
        # shuffle
        # NOTE: the dict of loader options may get used elsewhere!
        # However, this del doesn't touch those because loader_kwargs comes
        # from a **kwargs dict.
        del loader_kwargs["shuffle"]

    dataloader = DataLoader(dataset, **loader_kwargs)

    return dataloader


"""PyTorch compatible samplers.

https://github.com/speechbrain/speechbrain/blob/develop/speechbrain/dataio/sampler.py

These determine the order of iteration through a dataset.

Authors:
  * <NAME> 2020
  * <NAME> 2020
  * <NAME> 2020
"""
class ReproducibleRandomSampler(RandomSampler):
    """A modification of RandomSampler which always returns the same values.

    Also look at `torch.utils.data.RandomSampler`. This has mostly
    the same behaviour and arguments, except for adding 'seed' and 'epoch' and
    not supporting 'generator'.

    Note
    ----
    Call `set_epoch` before every epoch. Otherwise, the sampler will produce the
    same sequence of indices every epoch.

    Arguments
    ---------
    data_source : Dataset
        The data source to sample indices for.
    seed : int
        The base seed to use for the random number generator. It is recommended
        to use a value which has a good mix of 0 and 1 bits.
    epoch : int
        The epoch to start at.


    """

    def __init__(self, data_source, seed=563375142, epoch=0, **kwargs):
        if "generator" in kwargs:
            MSG = (
                "Cannot give a separate generator when using "
                + "ReproducibleRandomSampler"
            )
            raise ValueError(MSG)
        super().__init__(data_source, **kwargs)
        self.seed = int(seed)
        self.epoch = epoch
        self.generator = torch.Generator()

    def set_epoch(self, epoch):
        """
        You can also just access self.epoch, but we maintain this interface
        to mirror torch.utils.data.distributed.DistributedSampler
        """
        self.epoch = epoch

    def __iter__(self):
        self.generator.manual_seed(self.seed + self.epoch)
        return super().__iter__()



# https://github.com/speechbrain/speechbrain/blob/develop/speechbrain/utils/data_utils.py
def batch_pad_right(tensors: list, mode="constant", value=0):
    """Given a list of torch tensors it batches them together by padding to the right
    on each dimension in order to get same length for all.

    Parameters
    ----------
    tensors : list
        List of tensor we wish to pad together.
    mode : str
        Padding mode see torch.nn.functional.pad documentation.
    value : float
        Padding value see torch.nn.functional.pad documentation.

    Returns
    -------
    tensor : torch.Tensor
        Padded tensor.
    valid_vals : list
        List containing proportion for each dimension of original, non-padded values.

    """

    if not len(tensors):
        raise IndexError("Tensors list must not be empty")
    # tensors = list(map(list,tensors))

    if len(tensors) == 1:
        # if there is only one tensor in the batch we simply unsqueeze it.
        return tensors[0].unsqueeze(0), torch.tensor([1.0])

    if not (
        any(
            [tensors[i].ndim == tensors[0].ndim for i in range(1, len(tensors))]
        )
    ):
        raise IndexError("All tensors must have same number of dimensions")

    # FIXME we limit the support here: we allow padding of only the last dimension
    # need to remove this when feat extraction is updated to handle multichannel.
    max_shape = []
    for dim in range(tensors[0].ndim):
        if dim != (tensors[0].ndim - 1):
            if not all(
                [x.shape[dim] == tensors[0].shape[dim] for x in tensors[1:]]
            ):
                raise EnvironmentError(
                    "Tensors should have same dimensions except for last one"
                )
        max_shape.append(max([x.shape[dim] for x in tensors]))

    batched = []
    valid = []
    for t in tensors:
        # for each tensor we apply pad_right_to
        padded, valid_percent = pad_right_to(
            t, max_shape, mode=mode, value=value
        )
        batched.append(padded)
        valid.append(valid_percent[0])

    batched = torch.stack(batched)

    return batched, torch.tensor(valid)

def pad_right_to(
    tensor: torch.Tensor, target_shape: (list, tuple), mode="constant", value=0,
):
    """
    This function takes a torch tensor of arbitrary shape and pads it to target
    shape by appending values on the right.

    Parameters
    ----------
    tensor : input torch tensor
        Input tensor whose dimension we need to pad.
    target_shape : (list, tuple)
        Target shape we want for the target tensor its len must be equal to tensor.ndim
    mode : str
        Pad mode, please refer to torch.nn.functional.pad documentation.
    value : float
        Pad value, please refer to torch.nn.functional.pad documentation.

    Returns
    -------
    tensor : torch.Tensor
        Padded tensor.
    valid_vals : list
        List containing proportion for each dimension of original, non-padded values.
    """
    assert len(target_shape) == tensor.ndim
    pads = []  # this contains the abs length of the padding for each dimension.
    valid_vals = []  # thic contains the relative lengths for each dimension.
    i = len(target_shape) - 1  # iterating over target_shape ndims
    j = 0
    while i >= 0:
        assert (
            target_shape[i] >= tensor.shape[i]
        ), "Target shape must be >= original shape for every dim"
        pads.extend([0, target_shape[i] - tensor.shape[i]])
        valid_vals.append(tensor.shape[j] / target_shape[j])
        i -= 1
        j += 1

    tensor = torch.nn.functional.pad(tensor, pads, mode=mode, value=value)

    return tensor, valid_vals