import random
import torch
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.sampler import Sampler
from dataset.texts import phonemes_to_sequence
import numpy as np
from dataset.texts import text_to_sequence
from utils.util import pad_list, str_to_int_list, remove_outlier


def get_tts_dataset(path, batch_size, hp, valid=False):

    if valid:
        file_ = hp.data.valid_filelist
        pin_mem = False
        num_workers = 0
        shuffle = False
    else:
        file_ = hp.data.train_filelist
        pin_mem = True
        num_workers = 4
        shuffle = True
    train_dataset = TTSDataset(
        path, file_, hp.train.use_phonemes, hp.data.tts_cleaner_names, hp.train.eos
    )

    train_set = DataLoader(
        train_dataset,
        collate_fn=collate_tts,
        batch_size=batch_size,
        num_workers=num_workers,
        shuffle=shuffle,
        pin_memory=pin_mem,
    )
    return train_set


class TTSDataset(Dataset):
    def __init__(self, path, file_, use_phonemes, tts_cleaner_names, eos):
        self.path = path
        with open("{}".format(file_), encoding="utf-8") as f:
            self._metadata = [line.strip().split("|") for line in f]
        self.use_phonemes = use_phonemes
        self.tts_cleaner_names = tts_cleaner_names
        self.eos = eos

    def __getitem__(self, index):
        id = self._metadata[index][4].split(".")[0]
        x_ = self._metadata[index][3].split()
        if self.use_phonemes:
            x = phonemes_to_sequence(x_)
        else:
            x = text_to_sequence(x_, self.tts_cleaner_names, self.eos)
        mel = np.load(f"{self.path}mels/{id}.npy")
        durations = str_to_int_list(self._metadata[index][2])
        e = remove_outlier(
            np.load(f"{self.path}energy/{id}.npy")
        )  # self._norm_mean_std(np.load(f'{self.path}energy/{id}.npy'), self.e_mean, self.e_std, True)
        p = remove_outlier(
            np.load(f"{self.path}pitch/{id}.npy")
        )  # self._norm_mean_std(np.load(f'{self.path}pitch/{id}.npy'), self.f0_mean, self.f0_std, True)
        mel_len = mel.shape[1]
        durations = durations[: len(x)]
        durations[-1] = durations[-1] + (mel.shape[1] - sum(durations))
        assert mel.shape[1] == sum(durations)

        avg_mel = np.load(f"{self.path}avg_mel_ph/{id}.npy")
        assert avg_mel.shape[0] == len(x)


        return (
            np.array(x),
            mel.T,
            id,
            mel_len,
            np.array(durations),
            e,
            p,
            avg_mel,
        )  # Mel [T, num_mel]

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

    def _norm_mean_std(self, x, mean, std, is_remove_outlier=False):
        if is_remove_outlier:
            x = remove_outlier(x)
        zero_idxs = np.where(x == 0.0)[0]
        x = (x - mean) / std
        x[zero_idxs] = 0.0
        return x


def pad1d(x, max_len):
    return np.pad(x, (0, max_len - len(x)), mode="constant")


def pad2d(x, max_len):
    return np.pad(x, ((0, 0), (0, max_len - x.shape[-1])), mode="constant")


def collate_tts(batch):

    ilens = torch.from_numpy(np.array([x[0].shape[0] for x in batch])).long()
    olens = torch.from_numpy(np.array([y[1].shape[0] for y in batch])).long()
    ids = [x[2] for x in batch]

    # perform padding and conversion to tensor
    inputs = pad_list([torch.from_numpy(x[0]).long() for x in batch], 0)
    mels = pad_list([torch.from_numpy(y[1]).float() for y in batch], 0)

    durations = pad_list([torch.from_numpy(x[4]).long() for x in batch], 0)
    energys = pad_list([torch.from_numpy(y[5]).float() for y in batch], 0)
    pitches = pad_list([torch.from_numpy(y[6]).float() for y in batch], 0)

    avg_mel = pad_list([torch.from_numpy(y[7]).float() for y in batch], 0)

    # make labels for stop prediction
    labels = mels.new_zeros(mels.size(0), mels.size(1))
    for i, l in enumerate(olens):
        labels[i, l - 1 :] = 1.0

    # scale spectrograms to -4 <--> 4
    # mels = (mels * 8.) - 4

    return inputs, ilens, mels, labels, olens, ids, durations, energys, pitches, avg_mel


class BinnedLengthSampler(Sampler):
    def __init__(self, lengths, batch_size, bin_size):
        _, self.idx = torch.sort(torch.tensor(lengths).long())
        self.batch_size = batch_size
        self.bin_size = bin_size
        assert self.bin_size % self.batch_size == 0

    def __iter__(self):
        # Need to change to numpy since there's a bug in random.shuffle(tensor)
        # TODO : Post an issue on pytorch repo
        idx = self.idx.numpy()
        bins = []

        for i in range(len(idx) // self.bin_size):
            this_bin = idx[i * self.bin_size : (i + 1) * self.bin_size]
            random.shuffle(this_bin)
            bins += [this_bin]

        random.shuffle(bins)
        binned_idx = np.stack(bins).reshape(-1)

        if len(binned_idx) < len(idx):
            last_bin = idx[len(binned_idx) :]
            random.shuffle(last_bin)
            binned_idx = np.concatenate([binned_idx, last_bin])

        return iter(torch.tensor(binned_idx).long())

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