from tensorflow import keras
from tensorflow.keras import Model, layers
from tensorflow.keras.layers import Dense, Dropout, Conv2D
from tensorflow.keras.layers import LSTM, TimeDistributed, Bidirectional
from tensorflow.keras.constraints import max_norm
import librosa
import scipy
import numpy as np
import os
from ... import Metric

# prevent TF warnings
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'


class MOSNet(Metric):
    def __init__(self, window, hop=None):
        super(MOSNet, self).__init__(name='MOSNet', window=window, hop=hop)

        # constants
        self.fixed_rate = 16000
        self.mono = True
        self.absolute = True

        self.FFT_SIZE = 512
        self.SGRAM_DIM = self.FFT_SIZE // 2 + 1
        self.HOP_LENGTH = 256
        self.WIN_LENGTH = 512

        _input = keras.Input(shape=(None, 257))

        re_input = layers.Reshape((-1, 257, 1), input_shape=(-1, 257))(_input)

        # CNN
        conv1 = (Conv2D(16, (3, 3), strides=(1, 1), activation='relu',
                 padding='same'))(re_input)
        conv1 = (Conv2D(16, (3, 3), strides=(1, 1), activation='relu',
                 padding='same'))(conv1)
        conv1 = (Conv2D(16, (3, 3), strides=(1, 3), activation='relu',
                 padding='same'))(conv1)

        conv2 = (Conv2D(32, (3, 3), strides=(1, 1), activation='relu',
                 padding='same'))(conv1)
        conv2 = (Conv2D(32, (3, 3), strides=(1, 1), activation='relu',
                 padding='same'))(conv2)
        conv2 = (Conv2D(32, (3, 3), strides=(1, 3), activation='relu',
                 padding='same'))(conv2)

        conv3 = (Conv2D(64, (3, 3), strides=(1, 1), activation='relu',
                 padding='same'))(conv2)
        conv3 = (Conv2D(64, (3, 3), strides=(1, 1), activation='relu',
                 padding='same'))(conv3)
        conv3 = (Conv2D(64, (3, 3), strides=(1, 3), activation='relu',
                 padding='same'))(conv3)

        conv4 = (Conv2D(128, (3, 3), strides=(1, 1), activation='relu',
                 padding='same'))(conv3)
        conv4 = (Conv2D(128, (3, 3), strides=(1, 1), activation='relu',
                 padding='same'))(conv4)
        conv4 = (Conv2D(128, (3, 3), strides=(1, 3), activation='relu',
                 padding='same'))(conv4)

        re_shape = layers.Reshape((-1, 4*128), input_shape=(-1, 4, 128))(conv4)

        # BLSTM
        blstm1 = Bidirectional(
            LSTM(128, return_sequences=True, dropout=0.3,
                 recurrent_dropout=0.3,
                 recurrent_constraint=max_norm(0.00001)),
            merge_mode='concat')(re_shape)

        # DNN
        flatten = TimeDistributed(layers.Flatten())(blstm1)
        dense1 = TimeDistributed(Dense(128, activation='relu'))(flatten)
        dense1 = Dropout(0.3)(dense1)

        frame_score = TimeDistributed(Dense(1), name='frame')(dense1)
        import warnings

        average_score = layers.GlobalAveragePooling1D(name='avg')(frame_score)

        self.model = Model(outputs=[average_score, frame_score], inputs=_input)

        # weights are in the directory of this file
        pre_trained_dir = os.path.dirname(__file__)

        # load pre-trained weights. CNN_BLSTM is reported as best
        self.model.load_weights(os.path.join(pre_trained_dir, 'cnn_blstm.h5'))

    def test_window(self, audios, rate):
        # stft. D: (1+n_fft//2, T)
        linear = librosa.stft(y=np.asfortranarray(audios[0]),
                              n_fft=self.FFT_SIZE,
                              hop_length=self.HOP_LENGTH,
                              win_length=self.WIN_LENGTH,
                              window=scipy.signal.hamming,
                              )

        # magnitude spectrogram
        mag = np.abs(linear)  # (1+n_fft/2, T)

        # shape in (T, 1+n_fft/2)
        mag = np.transpose(mag.astype(np.float32))

        # now call the actual MOSnet
        return {'mosnet':
                self.model.predict(mag[None, ...], verbose=0, batch_size=1)[0]}