import numpy as np
from scipy.io.wavfile import read, write
import librosa


def mix_audio(signal_file, noise_file, snr):
    # if the audio is longer than the noise
    # play the noise in repeat for the duration of the audio
    signal_samples,s_rate= librosa.load(signal_file, sr=16000)
    noise_samples, s_rate= librosa.load(noise_file, sr=16000)
    print(signal_samples.shape, noise_samples.shape)

    signal=np.array(signal_samples,dtype=float)
    noise= np.array(noise_samples,dtype=float)

    noise = noise[np.arange(len(signal)) % len(noise)]
    
    # if the audio is shorter than the noi
    # this is important if loading resulted in 
    # uint8 or uint16 types, because it would cause overflow
    # when squaring and calculating mean
    noise = noise.astype(np.float32)
    signal = signal.astype(np.float32)
    
    # get the initial energy for reference
    signal_energy = np.mean(signal**2)
    noise_energy = np.mean(noise**2)
    # calculates the gain to be applied to the noise 
    # to achieve the given SNR
    g = np.sqrt(10.0 ** (-snr/10) * signal_energy / noise_energy)
    
    # Assumes signal and noise to be decorrelated
    # and calculate (a, b) such that energy of 
    # a*signal + b*noise matches the energy of the input signal
    a = np.sqrt(1 / (1 + g**2))
    b = np.sqrt(g**2 / (1 + g**2))
    print(g, a, b)

    monaural= a * signal + b * noise
    write('output.wav', 16000, monaural)


    # mix the signals
    return monaural