import numpy as np
import matplotlib.pyplot as plt
import scipy.signal
from scipy.io import wavfile
import sounddevice as sd
import random
from pesq import pesq



SIGNAL_PATH = "speechfiles/sp01.wav"
NOISE_PATH = "noisefiles/white.dat"


# Scale the signal to the range [-1,1]
def normalize_signal(signal):
    min_amp = np.min(signal)
    normalized_signal = signal - min_amp
    max_amp = np.max(normalized_signal)
    normalized_signal *= 2.0/max_amp
    normalized_signal -= 1
    return normalized_signal


# Load an audio file from disk
def load_audiofile(path):
    sound_data = []
    sample_rate = 8000
    # Load .dat files as sound files sampled at 8[kHz]
    if path[-3:] == "dat":
        with open(path, "r") as sound_file:
            sound_data_strings = sound_file.readlines()
            for data_string in sound_data_strings:
                sound_data.append(eval(data_string.strip()))
            sound_data = np.array(sound_data, dtype=np.float64)
    elif path[-3:] == "wav":
        sample_rate, sound_data = wavfile.read(path)
        # Make sure it is nparray of floats (trust me bro, normalizing yells at you if its ints)
        sound_data = np.array(sound_data, dtype=np.float64)
    return sample_rate, sound_data


# Add noise to a signal with a desired SNR
def add_noise(signal, noise, snr):
    len_signal = len(signal)
    len_noise = len(noise)

    # Get a random crop of the noise to match the length of the signal
    noise_crop_start = random.randrange(len_noise-len_signal)
    noise_crop = noise[noise_crop_start:noise_crop_start+len_signal]

    # Calculate the power of the signal and noise
    noise_power = np.linalg.norm(noise_crop, 2)
    signal_power = np.linalg.norm(signal, 2)

    # Adjust the noise level to match desired SNR
    u = 10**(snr/20)
    desired_noise_power = signal_power/u
    ratio = desired_noise_power / noise_power
    noise_crop *= ratio

    noisy_signal = signal + noise_crop
    return noisy_signal


def main():
    signal_sample_rate, signal_data = load_audiofile(SIGNAL_PATH)
    noise_sample_rate, noise_data = load_audiofile(NOISE_PATH)
    assert signal_sample_rate == noise_sample_rate, "Signal and noise sampling rates didn't match."
    sample_rate = signal_sample_rate

    noisy_signal = add_noise(signal_data, noise_data, 0)
    filtered_signal = scipy.signal.wiener(noisy_signal)

    print(pesq(sample_rate, signal_data, noisy_signal, mode='nb'))
    print(pesq(sample_rate, signal_data, filtered_signal, mode='nb'))
    #sd.play(normalize_signal(noisy_signal), samplerate=sample_rate, blocking=True)
    #sd.play(normalize_signal(filtered_signal), samplerate=sample_rate, blocking=True)


if __name__ == "__main__":
    main()