import librosa
import librosa.display
import os
import numpy as np
import matplotlib.pyplot as plt
from keras.models import model_from_json
from keras import optimizers
from keras_preprocessing.image import ImageDataGenerator

class classifier(object):
    def __init__(self, model, weight):
        json_file = open(model, 'r')
        loaded_model_json = json_file.read()
        json_file.close()
        self.loaded_model = model_from_json(loaded_model_json)
        self.loaded_model.load_weights(weight)
        self.datagen=ImageDataGenerator(rescale=1./255.,validation_split=0.25)
        self.loaded_model.compile(optimizers.rmsprop(lr=0.0005, decay=1e-6),loss="categorical_crossentropy",metrics=["accuracy"])
        self.loaded_model.summary()

    def predict(self,wav_to_predict):
        image = self.create_spectrogram(wav_to_predict)
        predict_generator = self.datagen.flow_from_directory(
                directory="images_to_predict/",
                batch_size=128,
                seed=42,
                shuffle=False,
                class_mode="categorical",
                target_size=(64,64))
        prediction = self.loaded_model.predict_generator(predict_generator,steps=1)
        os.remove(image)
        predicted_class_indices=np.argmax(prediction,axis=1)
        labels = {'anser': 0, 'columba': 1, 'hirundo': 2, 'passer': 3, 'sturnus': 4, 'turdus': 5, 'upupa': 6}
        labels = dict((v,k) for k,v in labels.items())
        predictions = [labels[k] for k in predicted_class_indices]
        print("The sound was a : ", predictions[0])
        return predictions
    
    def create_spectrogram(self, wav):
        plt.interactive(False)
        clip, sample_rate = librosa.load(wav, sr=None)
        fig = plt.figure(figsize=[0.72,0.72])
        ax = fig.add_subplot(111)
        ax.axes.get_xaxis().set_visible(False)
        ax.axes.get_yaxis().set_visible(False)
        ax.set_frame_on(False)
        spectogram = librosa.feature.melspectrogram(y=clip, sr=sample_rate)
        librosa.display.specshow(librosa.power_to_db(spectogram, ref=np.max))
        if not os.path.exists("images_to_predict/unknown/"):
            os.makedirs("images_to_predict/unknown/")
        file_name  = "images_to_predict/unknown/"+wav[:-4]+ '.jpg'
        plt.savefig(file_name, dpi=400, bbox_inches='tight',pad_inches=0)
        plt.close()    
        fig.clf()
        plt.close(fig)
        plt.close('all')
        return file_name