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
import json



#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():
    noise_paths = ("noisefiles/white.dat", "noisefiles/train.dat", "noisefiles/street.dat", "noisefiles/exhibition.dat")

    # Compose signal paths for the 30 sentences
    signal_paths = []
    for i in range(1,30+1):
        signal_paths.append(f"speechfiles/sp{i:02}.wav")
    
    # SNR in dB
    snrs = (0, 10, 20, 30)

    pesqs = {"unfiltered": {}, "filtered": {}}

    for snr in snrs:
        pesqs["unfiltered"][snr] = {}
        pesqs["filtered"][snr] = {}

        for noise_path in noise_paths:
            pesqs["filtered"][snr][noise_path[:-4]] = []
            pesqs["unfiltered"][snr][noise_path[:-4]] = []

            for signal_path in signal_paths:
                noise_sample_rate, noise_data = load_audiofile(noise_path)
                signal_sample_rate, signal_data = load_audiofile(signal_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, snr)
                filtered_signal = scipy.signal.wiener(noisy_signal)
            
                noisy_pesq = pesq(sample_rate, signal_data, noisy_signal, mode='nb')
                filtered_pesq = pesq(sample_rate, signal_data, filtered_signal, mode='nb')

                pesqs["filtered"][snr][noise_path[:-4]].append(noisy_pesq)
                pesqs["unfiltered"][snr][noise_path[:-4]].append(filtered_pesq)

    pesqs_json = json.dumps(pesqs, indent=4)
    with open("pesqs.json", "w") as outfile:
        outfile.write(pesqs_json)
    
    plt.plot(snrs, noisy_pesqs, label="PESQ Unfiltered")
    plt.plot(snrs, filtered_pesqs, label="PESQ Filtered")
    plt.xlabel("SNR [dB]")
    plt.ylabel("PESQ")
    plt.legend()
    plt.show()

    #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()