from keras.layers import Input, Add, Dense, Activation, ZeroPadding2D, BatchNormalization, Flatten, Conv2D, AveragePooling2D, MaxPooling2D, GlobalMaxPooling2D
from keras.initializers import glorot_uniform
from keras.models import Model, load_model
from custom_layers.scale_layer import Scale
from keras.optimizers import SGD

# Structure of an identity block:
from identity_block import *

# Structure of a Convolution Block:
from conv_block import *

def ResNet(input_shape = (224, 224, 3), classes = 196):
    """
    Implementation of the popular ResNet50 the following architecture:
    CONV2D -> BATCHNORM -> RELU -> MAXPOOL -> CONVBLOCK -> IDBLOCK*2 -> CONVBLOCK -> IDBLOCK*3
    -> CONVBLOCK -> IDBLOCK*5 -> CONVBLOCK -> IDBLOCK*2 -> AVGPOOL -> TOPLAYER

    Arguments:
    input_shape -- shape of the images of the dataset
    classes -- integer, number of classes

    Returns:
    model -- a Model() instance in Keras
    """
    # Set epsilon value for Batch Normalization
    eps = 1.1e-5

    # Define the input as a tensor with shape input_shape
    X_input = Input(input_shape)


    # Zero-Padding
    X = ZeroPadding2D((3, 3))(X_input)

    # Stage 1
    X = Conv2D(64, (7, 7), strides = (2, 2), name = 'conv1', use_bias=False, kernel_initializer = glorot_uniform(seed=0))(X)
    X = BatchNormalization(epsilon = eps, axis = 3, name = 'bn_conv1')(X)
    X = Scale(axis = 3, name = 'scale_conv1')(X)
    X = Activation('relu')(X)
    X = MaxPooling2D((3, 3), strides=(2, 2))(X)

    # Stage 2
    X = convolutional_block(X, f = 3, filters = [64, 64, 256], stage = 2, block='a', s = 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')

    # Stage 3 (Default Identity blocks = 3)
    X = convolutional_block(X, f = 3, filters = [128,128,512], stage = 3, block='a', s = 2)
    for i in range(1,8):
      X = identity_block(X, 3, [128, 128, 512], stage=3, block='b'+str(i))

    # Stage 4 (Default Identity blocks = 5)
    X = convolutional_block(X, f = 3, filters = [256, 256, 1024], stage = 4, block='a', s = 2)
    for i in range(1,36):
      X = identity_block(X, 3, [256, 256, 1024], stage=4, block='b'+str(i))


    # Stage 5
    X = convolutional_block(X, f = 3, filters = [512, 512, 2048], stage = 5, block='a', s = 2)
    X = identity_block(X, 3, [512, 512, 2048], stage=5, block='b')
    X = identity_block(X, 3, [512, 512, 2048], stage=5, block='c')

    # First output layer
    Xfc = AveragePooling2D(pool_size=(2, 2), padding='valid', name = "avg_pool")(X)
    Xfc = Flatten()(Xfc)
    Xfc = Dense(1000, activation='softmax', name='fc1000' + str(classes))(Xfc)

    # Transfer learning from pretrained weights
    model = Model(inputs = X_input, outputs = Xfc, name='Cifar10_transfer_learning')
    model.load_weights('models/resnet152_weights_tf.h5', by_name=True)

    # Switch to the final output layer
    Xfc2 = AveragePooling2D((7, 7), name='avg_pool')(X)
    Xfc2 = Flatten()(Xfc2)
    Xfc2 = Dense(classes, activation='softmax', name='fc'+str(classes))(Xfc2)


    # Create model
    model = Model(inputs = X_input, outputs = Xfc2, name='ResNet150')

    return model
