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


def pad(x, max_len=64600, random_clip=True):
    x_len = x.shape[0]
    if x_len > max_len:
        # random clip
        if random_clip:
            start_idx = np.random.randint(0, x_len - max_len)
            return x[start_idx:start_idx + max_len]
        else:
            return x[:max_len]
    # need to pad
    num_repeats = int(max_len / x_len)+1
    padded_x = np.tile(x, (1, num_repeats))[:, :max_len][0]
    return padded_x	

class DemoDataset(Dataset):
    def __init__(self, path):
        self.path = path
    
    def __len__(self):
        return 1
    
    def __getitem__(self, idx):
        waveform, sample_rate = librosa.load(self.path, sr=16000)
        waveform_pad = pad(waveform)
        waveform_tensor = Tensor(waveform_pad)

        return {
            'waveforms': waveform_tensor,
        }