import librosa
from tensorflow import keras
import numpy as np
import librosa.display
from matplotlib import pyplot as plt
from tqdm import tqdm
import os
import matplotlib.image as mpimg
import librosa
import soundfile as sf
import librosa.display


def reshape_npy_to_recurrent(npy_chapitre, longueur_original, longueur_sequence):
    """
    reshape the original array from  N,64,64,1 to N-4, 5, 64, 64 ,1

    for example
    10 images N = 10 10,64,64,1
    longeur sequence = 5
    reshape to 6,5,64,64,1
    12345
    23456
    34567
    45678
    56789
    678910     ->  10-5+1, 5, 64, 64, 1

    :param npy_chapitre: the numpy array with shape N,64,64,1
    :param longueur_original: N
    :param longueur_sequence: the number of images in one sequence
    :return: a np array with shape (longeur_original-longeur_sequence+1,longueur_sequence,64,64,1)
    """
    result = np.zeros((longueur_original-longueur_sequence+1, longueur_sequence, 64, 64, 1))
    for i in range(len(result)):
        result[i, :, :, :] = npy_chapitre[i:i+longueur_sequence, :, :, :]
    print(result.shape)
    return result


def compare_spectre(trained_model,
                    wav_file_griffin_lim,
                    lips_validation="../validation_data/lips_validation_ch7.npy",
                    tongues_validation="../validation_data/tongues_validation_ch7.npy",
                    spectre_validation="../validation_data/spectrum_validation.npy",
                    sample_rate=44100,
                    original_wav_file="../RecFile_1_20200617_153719_Sound_Capture_DShow_5_monoOutput1.wav"
                    ):
    """
    This function is used to compare the three parts of results using fft spectre :
    The original wav file, the spectrum generated, the wav file generated by griffin-lim
    :param lips_validation: the file path of the validation data lips
    :param tongues_validation: the file path of the validation data tongues
    :param spectre_validation: the file path of the validation data spectrum
    :param trained_model: the file path of the pre-trained model
    :param original_wav_file: the file path of the original wav file for the validation data
    :param wav_file_griffin_lim: the file path of the wav file reconstructed by griffin-lim
    :param sample_rate:
    :return:

    """
    # load validation data
    X_lips = np.load(lips_validation)
    X_tongues = np.load(tongues_validation)
    Y = np.load(spectre_validation)

    # normalisation
    X_lips = X_lips / 255.0
    X_tongues = X_tongues / 255.0
    max_spectrum = np.max(Y)
    Y = Y / max_spectrum
    y_test = np.transpose(Y)

    # load the pre-trained model
    model = keras.models.load_model(trained_model)
    model.summary()

    # calculate the spectrum predicted
    test_result = model.predict([X_lips, X_tongues])
    # calculate the fft predicted and normalize it between 0 and 1
    result = np.sum(test_result, axis=0)
    max_result = np.max(result)
    result /= max_result

    # calculate the stft of the original signal
    wav_original, _ = librosa.load(original_wav_file, sr=44100)
    spectrogramme_original = np.abs(librosa.stft(wav_original, n_fft=735 * 2, win_length=735 * 2, hop_length=735))
    # calculate the spectre and normalize it between 0 and 1
    spectre_original = np.transpose(spectrogramme_original)
    spectre_original = np.sum(spectre_original, axis=0)
    max_spectre_original = np.max(spectre_original)
    spectre_original /= max_spectre_original

    # calculate the stft of the reconstructed signal
    wav_reconstructed, _ = librosa.load(wav_file_griffin_lim)
    spectrogramme_wav_reconstructed = np.abs(librosa.stft(wav_reconstructed, n_fft=735 * 2,
                                                          win_length=735 * 2, hop_length=735))
    # calculate the spectre of the wav file reconstructed
    spectre_reconstructed = np.transpose(spectrogramme_wav_reconstructed)
    spectre_reconstructed = np.sum(spectre_reconstructed, axis=0)
    max_spectre_reconstructed = np.max(spectre_reconstructed)
    spectre_reconstructed /= max_spectre_original

    # plt the three spectre on a single figure
    plt.figure()
    x = np.arange(0, 736) * sample_rate / 1470
    plt.plot(x, spectre_original, 'r', label='spectre original')
    plt.plot(x, result, 'b', label='spectre predicted')
    plt.plot(x, spectre_reconstructed, 'g', label='spectre griffin lim')


if __name__ == "__main__":
    compare_spectre(trained_model="../results/ssi_model12-23-0.00004926.h5",
                    wav_file_griffin_lim="../results/ch7_0421_model12_4926e5.wav")
    # calculate the stft of the original signal
    wav_original, _ = librosa.load("../RecFile_1_20200617_153719_Sound_Capture_DShow_5_monoOutput1.wav", sr=44100)
    spectrogramme_original = np.abs(librosa.stft(wav_original, n_fft=735 * 2, win_length=735 * 2, hop_length=735))
    test_reconstruit = librosa.griffinlim(spectrogramme_original, n_iter=128, hop_length=735, win_length=735 * 2)
    spectrogramme_griffinlim_signal_original = np.abs(librosa.stft(test_reconstruit, n_fft=735 * 2,
                                                                   win_length=735 * 2, hop_length=735))
    spectre_griffinlim_signal_original = np.transpose(spectrogramme_griffinlim_signal_original)
    spectre_griffinlim_signal_original = np.sum(spectre_griffinlim_signal_original, axis=0)
    max_value = np.max(spectre_griffinlim_signal_original)
    spectre_griffinlim_signal_original /= max_value
    x = np.arange(0, 736) * 44100 / 1470
    plt.plot(x, spectre_griffinlim_signal_original, 'black', label="spectre_griffinlim_signal_original")
    plt.legend()
    plt.xlabel("Hz ")
    plt.ylabel("Amplitude")
    plt.grid()
    plt.show()
    plt.show()
