from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import numpy as np
import torch
import librosa


class AudioDataset(Dataset):
    def __init__(
        self,
        filepaths,
        labels,
        skip_times=None,
        num_classes=1,
        normalize="std",
        max_len=32000,
        random_sampling=True,
        train=False,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.filepaths = filepaths
        self.labels = labels
        self.skip_times = skip_times
        self.num_classes = num_classes
        self.random_sampling = random_sampling
        self.normalize = normalize
        self.max_len = max_len
        self.train = train
        if not self.train:
            assert (
                not self.random_sampling
            ), "Ensure random_sampling is disabled for val"

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

    def crop_or_pad(self, audio, max_len, random_sampling=True):
        audio_len = audio.shape[0]
        if random_sampling:
            diff_len = abs(max_len - audio_len)
            if audio_len < max_len:
                pad1 = np.random.randint(0, diff_len)
                pad2 = diff_len - pad1
                audio = np.pad(audio, (pad1, pad2), mode="constant")
            elif audio_len > max_len:
                idx = np.random.randint(0, diff_len)
                audio = audio[idx : (idx + max_len)]
        else:
            if audio_len < max_len:
                audio = np.pad(audio, (0, max_len - audio_len), mode="constant")
            elif audio_len > max_len:
                # Crop from the beginning
                # audio = audio[:max_len]

                # Crop from 3/4 of the audio
                # eq: l = (3x + t + x) => idx = 3x = (l - t) / 4 * 3
                idx = int((audio_len - max_len) / 4 * 3)
                audio = audio[idx : (idx + max_len)]
        return audio

    def __getitem__(self, idx):
        # Load audio
        audio, sr = librosa.load(self.filepaths[idx], sr=None)
        target = np.array([self.labels[idx]])

        # Trim start of audio (torchaudio.transforms.vad)
        if self.skip_times is not None:
            skip_time = self.skip_times[idx]
            audio = audio[int(skip_time*sr):]

        # Ensure fixed length
        audio = self.crop_or_pad(audio, self.max_len, self.random_sampling)

        if self.normalize == "std":
            audio /= np.maximum(np.std(audio), 1e-6)
        elif self.normalize == "minmax":
            audio -= np.min(audio)
            audio /= np.maximum(np.max(audio), 1e-6)

        audio = torch.from_numpy(audio).float()
        target = torch.from_numpy(target).float().squeeze()
        return {
            "audio": audio,
            "target": target,
        }


def get_dataloader(
    filepaths,
    labels,
    skip_times=None,
    batch_size=8,
    num_classes=1,
    max_len=32000,
    random_sampling=True,
    normalize="std",
    train=False,
    # drop_last=False,
    pin_memory=True,
    worker_init_fn=None,
    collate_fn=None,
    num_workers=0,
    distributed=False,
):
    dataset = AudioDataset(
        filepaths,
        labels,
        skip_times=skip_times,
        num_classes=num_classes,
        max_len=max_len,
        random_sampling=random_sampling,
        normalize=normalize,
        train=train,
    )

    if distributed:
        # drop_last is set to True to validate properly
        # Ref: https://discuss.pytorch.org/t/how-do-i-validate-with-pytorch-distributeddataparallel/172269/8
        sampler = torch.utils.data.distributed.DistributedSampler(
            dataset, shuffle=train, drop_last=not train
        )
    else:
        sampler = None

    dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=(sampler is None) and train,
        # drop_last=drop_last,
        num_workers=num_workers,
        pin_memory=pin_memory,
        worker_init_fn=worker_init_fn,
        collate_fn=collate_fn,
        sampler=sampler,
    )
    return dataloader