#!/usr/bin/env python3
from typing import Tuple
import tempfile
import os
import os.path
import shutil
import opentracing

import librosa
import librosa.display
import numpy
import matplotlib.pyplot
from tensorflow.keras.models import model_from_json
from tensorflow.keras import optimizers
from tensorflow.keras_preprocessing.image import ImageDataGenerator


class Classifier(object):

    def __init__(self, model_filename: str, weights_filename: str):
        with opentracing.tracer.start_active_span('classifier.loadModel'):
            with open(model_filename, 'r') as f:
                self.loaded_model = model_from_json(f.read())

        with opentracing.tracer.start_active_span('classifier.loadWeights'):
            self.loaded_model.load_weights(weights_filename)

        with opentracing.tracer.start_active_span('classifier.loadDatagen'):
            self.datagen = ImageDataGenerator(rescale=1. / 255., validation_split=0.25)

        with opentracing.tracer.start_active_span('classifier.compile'):
            self.loaded_model.compile(optimizers.RMSprop(lr=0.0005, decay=1e-6), loss="categorical_crossentropy",
                                      metrics=["accuracy"])

        self.loaded_model.summary() # prints to stdout??

    @staticmethod
    def create_spectrogram(wav_filename: str) -> Tuple[str, str]:

        with opentracing.tracer.start_active_span('classifier.librosa.load'):
            clip, sample_rate = librosa.load(wav_filename, sr=None)

        with opentracing.tracer.start_active_span('classifier.Plot'):
            matplotlib.pyplot.interactive(False)
            fig = matplotlib.pyplot.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=numpy.max))

            target_dir = tempfile.mkdtemp(dir="/dev/shm")
            wav_basename = os.path.basename(wav_filename)

            # Change extension to jpg... mert 110% biztos vagyok benne hogy a keras nem bírná beolvasni máshogy
            file_name = os.path.join(target_dir, "unknown", f"{wav_basename[:-4]}.jpg")
            os.mkdir(os.path.join(target_dir, "unknown"))

            matplotlib.pyplot.savefig(file_name, dpi=400, bbox_inches='tight', pad_inches=0)
            matplotlib.pyplot.close()
            fig.clf()
            matplotlib.pyplot.close(fig)
            matplotlib.pyplot.close('all')

        return target_dir, file_name  # Az unknown nélkülivel kell visszatérni

    def _run_predictor(self, directory: str) -> Tuple[str, dict]:
        predict_generator = self.datagen.flow_from_directory(
            directory=directory,
            batch_size=128,
            seed=42,
            shuffle=False,
            class_mode="categorical",
            target_size=(64, 64))

        prediction = self.loaded_model.predict_generator(predict_generator, steps=1)

        predicted_class_indices = numpy.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())

        labeled_predictions = {labels[i]: float(p) for i, p in enumerate(prediction[0])}

        predicted_class_name = [labels[k] for k in predicted_class_indices][0]  # eh?

        return predicted_class_name, labeled_predictions

    def predict(self, wav_filename: str) -> Tuple[str, dict]:
        with opentracing.tracer.start_active_span('classifier.createSpectrogram'):
            directory, _ = self.create_spectrogram(wav_filename)

        with opentracing.tracer.start_active_span('classifier.runPredictor'):
            result = self._run_predictor(directory)

        shutil.rmtree(directory)  # The image is no longer needed

        return result