import torch


def get_samples_from_datasets(datasets, wav):
    """Gets samples (noise or speech) from the datasets.

    Arguments
    ---------
    datasets : list
        List containing datasets. More precisely, we expect here the pointers
        to the object used in speechbrain for data augmentation
        (e.g, speechbrain.lobes.augment.EnvCorrupt).
    wav : torch.tensor
        The original waveform. The drawn samples will have the same
        dimensionality of the original waveform.

    Returns
    -------
    samples: torch.tensor
        A batch of new samples drawn from the input list of datasets.
    """
    # We want a sample of the same size of the original signal
    samples = torch.zeros(
        wav.shape[0], wav.shape[1], len(datasets), device=wav.device
    )

    # Let's sample a sequence from each dataset
    for i, dataset in enumerate(datasets):

        # Initialize the signal with noise
        wav_sample = (torch.rand_like(wav) * 2) - 1
        len_sample = torch.ones(wav.shape[0], device=wav.device)

        # Sample a sequence
        wav_sample = dataset(wav_sample, len_sample)

        # Append it
        samples[:, :, i] = wav_sample

    # Random permutations of the signal
    idx = torch.randperm(samples.shape[-1], device=wav.device)
    samples[:, :] = samples[:, :, idx]
    return samples