# -*- coding: utf-8 -*-
"""
Created on Thu Feb  1 17:00:20 2018

@author: damodara
"""

# -*- coding: utf-8 -*-
"""
Created on Mon Oct  2 11:04:48 2017

@author: damodara
"""
import dnn
from keras.applications.vgg16 import VGG16
from keras.utils.vis_utils import model_to_dot
from keras.applications.vgg16 import preprocess_input
from keras.layers import Dense, GlobalAveragePooling2D, Flatten
import keras


    
def softmax_classification(input_shape, nclass, l2_weight=0):
    model = dnn.Sequential()
    model.add(dnn.Dense(nclass, input_shape=(input_shape,), activation='softmax',
                        kernel_regularizer=dnn.keras.regularizers.l2(l2_weight)))
    return model
    

    
    
    


def mnist_feat_ext(main_input, l2_weight=0.0):
    net = dnn.Convolution2D(32, (5, 5), padding='same', activation='relu',
                                kernel_regularizer=dnn.keras.regularizers.l2(l2_weight))(main_input)
    net = dnn.MaxPooling2D(pool_size=(2,2), strides=2)(net)  
    net = dnn.Convolution2D(48, (5,5),padding='same', activation='relu',
                                kernel_regularizer=dnn.keras.regularizers.l2(l2_weight))(net)
    net = dnn.MaxPooling2D(pool_size=(2,2),strides=2)(net)      
    net = dnn.Flatten()(net)
    # net = dnn.BatchNormalization(axis=1)(net)
    net = dnn.Dense(100,activation='relu',
                        kernel_regularizer=dnn.keras.regularizers.l2(l2_weight))(net)
    net = dnn.Dense(100,activation='tanh',
                        kernel_regularizer=dnn.keras.regularizers.l2(l2_weight))(net)
    return net


    
def assda_feat_ext(main_input, l2_weight=0.0, small_model=False):
    padding = 'same'
    maxpool_strides = 2 if small_model else 1
    net = dnn.Convolution2D(32, (3, 3),padding=padding, activation='relu',
                                kernel_regularizer=dnn.keras.regularizers.l2(l2_weight))(main_input)
    net = dnn.Convolution2D(32, (3, 3), padding=padding, activation='relu',
                                kernel_regularizer=dnn.keras.regularizers.l2(l2_weight))(net)
    net = dnn.MaxPooling2D(pool_size=(2, 2), strides=maxpool_strides)(net)
    # net = dnn.Dropout(0.5)(net)
    net = dnn.Convolution2D(64, (3, 3), padding=padding, activation='relu',
                                kernel_regularizer=dnn.keras.regularizers.l2(l2_weight))(net)
    net = dnn.Convolution2D(64, (3, 3), padding=padding, activation='relu',
                                kernel_regularizer=dnn.keras.regularizers.l2(l2_weight))(net)
    net = dnn.MaxPooling2D(pool_size=(2, 2), strides=maxpool_strides)(net)
    # net = dnn.Dropout(0.5)(net)
    net = dnn.Convolution2D(128, (3, 3), padding=padding, activation='relu',
                                kernel_regularizer=dnn.keras.regularizers.l2(l2_weight))(net)
    net = dnn.Convolution2D(128, (3, 3), padding=padding, activation='relu',
                                kernel_regularizer=dnn.keras.regularizers.l2(l2_weight))(net)
    net = dnn.MaxPooling2D(pool_size=(2, 2), strides=maxpool_strides)(net)
#   
    net = dnn.Flatten()(net)
    net = dnn.Dense(128,activation='sigmoid',
                        kernel_regularizer=dnn.keras.regularizers.l2(l2_weight),name='feat_ext')(net)
    return net
    

    
def regressor(model_input, noutputs, l2_weight=0.0):
    net = dnn.Dense(noutputs, activation='sigmoid', name='regressor_output')(model_input)
    return net


def classifier(model_input, nclass,l2_weight=0.0):
    net = dnn.Dense(nclass,activation='softmax', name='classifier_output')(model_input)
    return net
    
def classifier_dropout(model_input, nclass,l2_weight=0.0):
    net = dnn.Dropout(0.5)(model_input)
    net = dnn.Dense(nclass,activation='softmax', name='classifier_output')(net)
    return net
    


def res_net50_fe(img_input,l2_weight=0.0):
    from keras.models import Sequential
    from keras.optimizers import SGD, Adam
    from keras.layers import Input, Dense, Convolution2D, MaxPooling2D, AveragePooling2D, ZeroPadding2D, Dropout, \
        Flatten, merge, Reshape, Activation
    from keras.layers.normalization import BatchNormalization
    from keras.models import Model
    from keras import backend as K

    from sklearn.metrics import log_loss



    def identity_block(input_tensor, kernel_size, filters, stage, block):
        """
        The identity_block is the block that has no conv layer at shortcut
        Arguments
            input_tensor: input tensor
            kernel_size: defualt 3, the kernel size of middle conv layer at main path
            filters: list of integers, the nb_filters of 3 conv layer at main path
            stage: integer, current stage label, used for generating layer names
            block: 'a','b'..., current block label, used for generating layer names
        """

        nb_filter1, nb_filter2, nb_filter3 = filters
        conv_name_base = 'res' + str(stage) + block + '_branch'
        bn_name_base = 'bn' + str(stage) + block + '_branch'

        x = Convolution2D(nb_filter1, 1, 1, name=conv_name_base + '2a')(input_tensor)
        x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2a')(x)
        x = Activation('relu')(x)

        x = Convolution2D(nb_filter2, kernel_size, kernel_size,
                          border_mode='same', name=conv_name_base + '2b')(x)
        x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2b')(x)
        x = Activation('relu')(x)

        x = Convolution2D(nb_filter3, 1, 1, name=conv_name_base + '2c')(x)
        x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2c')(x)

        x = merge([x, input_tensor], mode='sum')
        x = Activation('relu')(x)
        return x

    def conv_block(input_tensor, kernel_size, filters, stage, block, strides=(2, 2)):
        """
        conv_block is the block that has a conv layer at shortcut
        # Arguments
            input_tensor: input tensor
            kernel_size: defualt 3, the kernel size of middle conv layer at main path
            filters: list of integers, the nb_filters of 3 conv layer at main path
            stage: integer, current stage label, used for generating layer names
            block: 'a','b'..., current block label, used for generating layer names
        Note that from stage 3, the first conv layer at main path is with subsample=(2,2)
        And the shortcut should have subsample=(2,2) as well
        """

        nb_filter1, nb_filter2, nb_filter3 = filters
        conv_name_base = 'res' + str(stage) + block + '_branch'
        bn_name_base = 'bn' + str(stage) + block + '_branch'

        x = Convolution2D(nb_filter1, 1, 1, subsample=strides,
                          name=conv_name_base + '2a')(input_tensor)
        x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2a')(x)
        x = Activation('relu')(x)

        x = Convolution2D(nb_filter2, kernel_size, kernel_size, border_mode='same',
                          name=conv_name_base + '2b')(x)
        x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2b')(x)
        x = Activation('relu')(x)

        x = Convolution2D(nb_filter3, 1, 1, name=conv_name_base + '2c')(x)
        x = BatchNormalization(axis=bn_axis, name=bn_name_base + '2c')(x)

        shortcut = Convolution2D(nb_filter3, 1, 1, subsample=strides,
                                 name=conv_name_base + '1')(input_tensor)
        shortcut = BatchNormalization(axis=bn_axis, name=bn_name_base + '1')(shortcut)

        x = merge([x, shortcut], mode='sum')
        x = Activation('relu')(x)
        return x

    def mean_subtract(img):
        img = dnn.Lambda(lambda x: x - 123.68)(img[:,:,0])
        img = dnn.Lambda(lambda x: x - 116.779)(img[:,:,1])
        img = dnn.Lambda(lambda x: x - 103.939) (img[:,:,2])

        return img / 255.0

    def resnet50_model(img_input, num_classes=None):
        """
        Resnet 50 Model for Keras

        Model Schema is based on 
        https://github.com/fchollet/deep-learning-models/blob/master/resnet50.py

        ImageNet Pretrained Weights 
        https://github.com/fchollet/deep-learning-models/releases/download/v0.2/resnet50_weights_th_dim_ordering_th_kernels.h5

        Parameters:
          img_rows, img_cols - resolution of inputs
          channel - 1 for grayscale, 3 for color 
          num_classes - number of class labels for our classification task
        """

        # Handle Dimension Ordering for different backends
        global bn_axis
        if K.image_dim_ordering() == 'tf':
            bn_axis = 3
            #img_input = Input(shape=(img_rows, img_cols, color_type))
        else:
            bn_axis = 1
            #img_input = Input(shape=(color_type, img_rows, img_cols))
        # img_input = keras.layers.Lambda(mean_subtract, name='mean_subtraction')(img_input)
        x = ZeroPadding2D((3, 3))(img_input)
        x = Convolution2D(64, 7, 7, subsample=(2, 2), name='conv1')(x)
        x = BatchNormalization(axis=bn_axis, name='bn_conv1')(x)
        x = Activation('relu')(x)
        x = MaxPooling2D((3, 3), strides=(2, 2))(x)

        x = conv_block(x, 3, [64, 64, 256], stage=2, block='a', strides=(1, 1))
        x = identity_block(x, 3, [64, 64, 256], stage=2, block='b')
        x = identity_block(x, 3, [64, 64, 256], stage=2, block='c')

        x = conv_block(x, 3, [128, 128, 512], stage=3, block='a')
        x = identity_block(x, 3, [128, 128, 512], stage=3, block='b')
        x = identity_block(x, 3, [128, 128, 512], stage=3, block='c')
        x = identity_block(x, 3, [128, 128, 512], stage=3, block='d')

        x = conv_block(x, 3, [256, 256, 1024], stage=4, block='a')
        x = identity_block(x, 3, [256, 256, 1024], stage=4, block='b')
        x = identity_block(x, 3, [256, 256, 1024], stage=4, block='c')
        x = identity_block(x, 3, [256, 256, 1024], stage=4, block='d')
        x = identity_block(x, 3, [256, 256, 1024], stage=4, block='e')
        x = identity_block(x, 3, [256, 256, 1024], stage=4, block='f')

        x = conv_block(x, 3, [512, 512, 2048], stage=5, block='a')
        x = identity_block(x, 3, [512, 512, 2048], stage=5, block='b')
        x = identity_block(x, 3, [512, 512, 2048], stage=5, block='c')

        # Fully Connected Softmax Layer
        x_fc = AveragePooling2D((7, 7), name='avg_pool')(x)
        x_fc = Flatten()(x_fc)
        x_fc = Dense(1000, activation='softmax', name='fc1000')(x_fc)

        # Create model
        model = Model(img_input, x_fc)

        # Load ImageNet pre-trained data
        if K.image_dim_ordering() == 'th':
            # Use pre-trained weights for Theano backend
            weights_path = '/home/damodara/DeepNetModels/cnn_finetune/imagenet_models/resnet50_weights_th_dim_ordering_th_kernels.h5'
        else:
            # Use pre-trained weights for Tensorflow backend
            weights_path = '/home/damodara/DeepNetModels/cnn_finetune/imagenet_models/resnet50_weights_tf_dim_ordering_tf_kernels.h5'

        model.load_weights(weights_path)

        # Truncate and replace softmax layer for transfer learning
        # Cannot use model.layers.pop() since model is not of Sequential() type
        # The method below works since pre-trained weights are stored in layers but not in the model
        x_newfc = AveragePooling2D((7, 7), name='avg_pool')(x)
        x_newfc = Flatten()(x_newfc)
        x_newfc = Dense(256, activation='relu', name='feat_ext')(x_newfc)
        #x_newfc = Dense(num_classes, activation='softmax', name='fc10')(x_newfc)

        # Create another model with our customized softmax
        #model = Model(img_input, x_newfc)

        # Learning rate is changed to 0.001
        #    sgd = SGD(lr=1e-3, decay=1e-6, momentum=0.9, nesterov=True)

        return x_newfc

    x_newfc = resnet50_model(img_input=img_input)
    return  x_newfc




def vgg16_fe(img_input):
    # net = preprocess_input(img_input)
    vgg_model = VGG16(weights='imagenet', include_top=True, input_tensor=img_input)
    vgg_model.layers.pop()
    return vgg_model.layers[-1].output
    # return  model.layers[-1].output

def vgg16F_fe(img_input):
    # net = preprocess_input(img_input)
    from keras_vggface.vggface import VGGFace
    vgg_model = VGGFace(include_top=False, input_tensor=img_input, pooling='avg')
    #vgg_model.layers.pop()
    last_layer = vgg_model.get_layer('pool5').output
    x = Flatten(name='flatten')(last_layer)
    x = Dense(1024, activation='relu', trainable=True)(x)
    x = Dense(512, activation='relu', trainable=True)(x)
    model = dnn.Model(input=vgg_model.input, output=x)
    return model.layers[-1].output






