import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from tensorflow.keras.models import Sequential, Model
from tensorflow.keras.layers import Dense, Input
from tensorflow.keras.layers import Dropout, GlobalMaxPooling2D
from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization ,Reshape
from tensorflow.keras.optimizers import SGD
from tensorflow.keras.layers import Activation, Layer
from tensorflow.keras.initializers import GlorotUniform
import tensorflow.keras.backend as K

class SelfAttention(Layer):
    def __init__(self, 
                 n_hop,
                 hidden_dim,
                 nc=256,
                 penalty=1.0,
                 return_attention=False,
                 kernel_initializer=GlorotUniform(),
                 kernel_regularizer=None,
                 kernel_constraint=None,
                 **kwargs):
        self.n_hop = n_hop
        self.hidden_dim = hidden_dim
        self.nc=nc
        self.penalty = penalty
        self.kernel_initializer = GlorotUniform() # tf.keras.initializers.get(kernel_initializer)
        self.kernel_regularizer = None #tf.keras.regularizers.get(kernel_regularizer)
        self.kernel_constraint = None #tf.keras.constraints.get(kernel_constraint)
        self.return_attention = return_attention
        super(SelfAttention, self).__init__(**kwargs)

    def build(self, input_shape):
        # input_shape: (None, Sequence_size, Sequence_hidden_dim)
        assert len(input_shape) >= 3
        batch_size, T, nh = input_shape
        
        self.Ws1 = self.add_weight(shape=(self.hidden_dim, self.nc),
                                      initializer=self.kernel_initializer,
                                      name='SelfAttention-Ws1',
                                      regularizer=self.kernel_regularizer,
                                      constraint=self.kernel_constraint)
        
        self.Ws2 = self.add_weight(shape=(self.nc, self.n_hop), 
                                      initializer=self.kernel_initializer,
                                      name='SelfAttention-Ws2',
                                      regularizer=self.kernel_regularizer,
                                      constraint=self.kernel_constraint)
        
        super(SelfAttention, self).build(input_shape)

    def compute_output_shape(self, input_shape):
        assert input_shape and len(input_shape) >= 3
        assert input_shape[-1]
        batch_size, sequence_size, sequence_hidden_dim = input_shape
        output_shape = tuple([batch_size, self.n_hop, sequence_hidden_dim])
        
        if self.return_attention:
            attention_shape = tuple([batch_size, self.n_hop, sequence_size])
            return [output_shape, attention_shape]
        else: return output_shape



    
    def _frobenius_norm(self, inputs):
        outputs = K.sqrt(K.sum(K.square(inputs)))
        return outputs    

    def call(self, inputs):
        shape=inputs.shape
        H=inputs
        x = K.tanh(tf.matmul(H,self.Ws1))
        x = tf.matmul(x,self.Ws2)
        A = K.softmax(x,axis=0) # A = softmax(dot(Ws2, d1))
        At=K.permute_dimensions(A,(0,2,1))
        E = tf.matmul(At,H)
        
        return E
   
    def get_config(self):

        config = super().get_config().copy()
        config.update({
            'n_hop': self.n_hop,
            'hidden_dim': self.hidden_dim,
            'nc': self.nc,
            'penalty': self.penalty,
            'kernel_initializer': self.kernel_initializer,
            'kernel_regularizer': self.kernel_regularizer,
            'kernel_constraint': self.kernel_constraint,
            'return_attention': self.return_attention,
        })
        return config


def vgg_att(n_class):
    inputs = Input(shape=(300,40,1))
    x=Conv2D(64, (3, 3), padding='same', name='block1_conv1',activation='relu')(inputs)
    x=Conv2D(64, (3, 3), padding='same', name='block1_conv2',activation='relu')(x)
    x=MaxPooling2D(pool_size = (2, 2), strides = (2, 2))(x)
    x=BatchNormalization()(x)
    x=Dropout(0.2)(x)


    print(x.shape)

    x=Conv2D(128, (3, 3), padding='same', name='block2_conv1',activation='relu')(x)
    x=Conv2D(128, (3, 3), padding='same', name='block2_conv2',activation='relu')(x)
    x=MaxPooling2D(pool_size = (2, 2), strides = (2, 2))(x)
    x=BatchNormalization()(x)
    x=Dropout(0.2)(x)
    print(x.shape)


    x=Conv2D(256, (3, 3), padding='same', name='block3_conv1',activation='relu')(x)
    x=Conv2D(256, (3, 3), padding='same', name='block3_conv2',activation='relu')(x)
    x=MaxPooling2D(pool_size = (2, 2), strides = (2, 2),padding="same")(x)
    x=BatchNormalization()(x)
    x=Dropout(0.2)(x)
    print(x.shape)

    x=Conv2D(512, (3, 3), padding='same', name='block4_conv1',activation='relu')(x)
    x=Conv2D(512, (3, 3), padding='same', name='block4_conv2',activation='relu')(x)
    x=MaxPooling2D(pool_size = (2, 2), strides = (2, 2),padding="same")(x)
    x=BatchNormalization()(x)
    x=Dropout(0.2)(x)
    print(x.shape)

    att=SelfAttention(n_hop=4,hidden_dim=1536)
    x=Reshape((x.shape[1], x.shape[2]*x.shape[3]))(x)
    print("after reshape")
    print(x.shape)
    x=att(x)
    print("after attention")
    print(x.shape)
    x=AveragePooling1D(pool_size=4,data_format="channels_last")(x)
    #x = GlobalMaxPooling2D()(x)
    print("after avgpool")
    print(x.shape)
    x = Flatten()(x)
    x = Dense(256, activation = 'relu')(x)
    x=Dropout(0.4)(x)
    output = Dense(n_class,activation = 'softmax')(x)
    model = Model(inputs=inputs, outputs=output)

    model.compile(loss='categorical_crossentropy',optimizer ='adam')#need hyperparam-tuning 
    model.summary()
    return model









tensorflow.keras.layers.AveragePooling1D tensorflow.matmul tensorflow.keras.backend.softmax tensorflow.keras.backend.permute_dimensions tensorflow.keras.models.Model tensorflow.keras.layers.Dense tensorflow.keras.layers.Conv2D tensorflow.keras.layers.Reshape tensorflow.keras.layers.MaxPooling2D tensorflow.keras.backend.square tensorflow.keras.layers.BatchNormalization tensorflow.keras.layers.Dropout tensorflow.keras.initializers.GlorotUniform tensorflow.keras.layers.Flatten tensorflow.keras.layers.Input
[(99, 'tensorflow.keras.layers.Input', 'Input', ([], {'shape': '(300, 40, 1)'}), False, 'from tensorflow.keras.layers import Dense, Input\n'), (146, 'tensorflow.keras.models.Model', 'Model', ([], {'inputs': 'inputs', 'outputs': 'output'}), False, 'from tensorflow.keras.models import Sequential, Model\n'), (20, 'tensorflow.keras.initializers.GlorotUniform', 'GlorotUniform', ([], {}), False, 'from tensorflow.keras.initializers import GlorotUniform\n'), (28, 'tensorflow.keras.initializers.GlorotUniform', 'GlorotUniform', ([], {}), False, 'from tensorflow.keras.initializers import GlorotUniform\n'), (75, 'tensorflow.matmul', 'tf.matmul', (['x', 'self.Ws2'], {}), True, 'import tensorflow as tf\n'), (76, 'tensorflow.keras.backend.softmax', 'K.softmax', (['x'], {'axis': '(0)'}), True, 'import tensorflow.keras.backend as K\n'), (77, 'tensorflow.keras.backend.permute_dimensions', 'K.permute_dimensions', (['A', '(0, 2, 1)'], {}), True, 'import tensorflow.keras.backend as K\n'), (78, 'tensorflow.matmul', 'tf.matmul', (['At', 'H'], {}), True, 'import tensorflow as tf\n'), (100, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(64)', '(3, 3)'], {'padding': '"""same"""', 'name': '"""block1_conv1"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (101, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(64)', '(3, 3)'], {'padding': '"""same"""', 'name': '"""block1_conv2"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (102, 'tensorflow.keras.layers.MaxPooling2D', 'MaxPooling2D', ([], {'pool_size': '(2, 2)', 'strides': '(2, 2)'}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (103, 'tensorflow.keras.layers.BatchNormalization', 'BatchNormalization', ([], {}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (104, 'tensorflow.keras.layers.Dropout', 'Dropout', (['(0.2)'], {}), False, 'from tensorflow.keras.layers import Dropout, GlobalMaxPooling2D\n'), (109, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(128)', '(3, 3)'], {'padding': '"""same"""', 'name': '"""block2_conv1"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (110, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(128)', '(3, 3)'], {'padding': '"""same"""', 'name': '"""block2_conv2"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (111, 'tensorflow.keras.layers.MaxPooling2D', 'MaxPooling2D', ([], {'pool_size': '(2, 2)', 'strides': '(2, 2)'}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (112, 'tensorflow.keras.layers.BatchNormalization', 'BatchNormalization', ([], {}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (113, 'tensorflow.keras.layers.Dropout', 'Dropout', (['(0.2)'], {}), False, 'from tensorflow.keras.layers import Dropout, GlobalMaxPooling2D\n'), (117, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(256)', '(3, 3)'], {'padding': '"""same"""', 'name': '"""block3_conv1"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (118, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(256)', '(3, 3)'], {'padding': '"""same"""', 'name': '"""block3_conv2"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (119, 'tensorflow.keras.layers.MaxPooling2D', 'MaxPooling2D', ([], {'pool_size': '(2, 2)', 'strides': '(2, 2)', 'padding': '"""same"""'}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (120, 'tensorflow.keras.layers.BatchNormalization', 'BatchNormalization', ([], {}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (121, 'tensorflow.keras.layers.Dropout', 'Dropout', (['(0.2)'], {}), False, 'from tensorflow.keras.layers import Dropout, GlobalMaxPooling2D\n'), (124, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(512)', '(3, 3)'], {'padding': '"""same"""', 'name': '"""block4_conv1"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (125, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(512)', '(3, 3)'], {'padding': '"""same"""', 'name': '"""block4_conv2"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (126, 'tensorflow.keras.layers.MaxPooling2D', 'MaxPooling2D', ([], {'pool_size': '(2, 2)', 'strides': '(2, 2)', 'padding': '"""same"""'}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (127, 'tensorflow.keras.layers.BatchNormalization', 'BatchNormalization', ([], {}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (128, 'tensorflow.keras.layers.Dropout', 'Dropout', (['(0.2)'], {}), False, 'from tensorflow.keras.layers import Dropout, GlobalMaxPooling2D\n'), (132, 'tensorflow.keras.layers.Reshape', 'Reshape', (['(x.shape[1], x.shape[2] * x.shape[3])'], {}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (138, 'tensorflow.keras.layers.AveragePooling1D', 'AveragePooling1D', ([], {'pool_size': '(4)', 'data_format': '"""channels_last"""'}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (142, 'tensorflow.keras.layers.Flatten', 'Flatten', ([], {}), False, 'from tensorflow.keras.layers import Flatten, Conv2D, MaxPooling2D, ZeroPadding2D, AveragePooling1D, BatchNormalization, Reshape\n'), (143, 'tensorflow.keras.layers.Dense', 'Dense', (['(256)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Dense, Input\n'), (144, 'tensorflow.keras.layers.Dropout', 'Dropout', (['(0.4)'], {}), False, 'from tensorflow.keras.layers import Dropout, GlobalMaxPooling2D\n'), (145, 'tensorflow.keras.layers.Dense', 'Dense', (['n_class'], {'activation': '"""softmax"""'}), False, 'from tensorflow.keras.layers import Dense, Input\n'), (74, 'tensorflow.matmul', 'tf.matmul', (['H', 'self.Ws1'], {}), True, 'import tensorflow as tf\n'), (68, 'tensorflow.keras.backend.square', 'K.square', (['inputs'], {}), True, 'import tensorflow.keras.backend as K\n')]