import torch
import random
from absl import flags
import scipy.io.wavfile as wav

def create_speech_rir(audio, rir):
    """
    Returns:
        A tensor of speech with reverberations (Convolve the audio with the rir)
    """
    shape = audio.shape[0]

    # Compute convolution in fourier space
    sp1 = torch.fft.rfft(rir, shape)
    sp2 = torch.fft.rfft(audio, shape)
    ret = torch.fft.irfft(sp1 * sp2, shape)

    # Normalization
    ret = ret / torch.abs(ret).max()

    return ret

def Readrir(root_path, num_rir, eval=False):
    '''
    Return:
        rir: a numpy array of the room reverberation

    '''
    if not eval:
        index = random.randint(1, num_rir)
        _, rir = wav.read(root_path + "/3575-170457-0013_rir_" + str(index) + ".wav")
    else:
        index = random.randint(1000, 1200)
        _, rir = wav.read(root_path + "/3575-170457-0013_rir_" + str(index) + ".wav")
    return rir