from keras.layers import Input, merge, Dropout, Dense, Flatten, Activation
from keras.layers.convolutional import MaxPooling2D, Convolution2D
from keras.layers.covonlutional import AveragePooling2D
from keras.layers.normalization import BatchNormalization
from keras.models import Model
from keras.callbacks import EarlyStopping, ReduceLROnPlateau

from keras import backend as K
from preprocess import process_test, process_train, process_validation

"""Implementation of Inception Network v4 [Inception Network v4
Paper](http://arxiv.org/pdf/1602.07261v1.pdf) in Keras.

"""


def inception_stem(input):  # Input (299,299,3)
    if K.image_dim_ordering() == "th":
        channel_axis = 1
    else:
        channel_axis = -1

    # Input Shape is 299 x 299 x 3 (th) or 3 x 299 x 299 (th)
    c = Convolution2D(32, 3, 3, activation='relu', subsample=(2, 2))(input)
    c = Convolution2D(32, 3, 3, activation='relu', )(c)
    c = BatchNormalization(axis=channel_axis)(c)
    c = Convolution2D(64, 3, 3, activation='relu', border_mode='same')(c)

    c1 = MaxPooling2D((3, 3), strides=(2, 2))(c)
    c2 = Convolution2D(96, 3, 3, activation='relu', subsample=(2, 2))(c)

    m = merge([c1, c2], mode='concat', concat_axis=channel_axis)

    c1 = Convolution2D(64, 1, 1, activation='relu', border_mode='same')(m)
    c1 = Convolution2D(96, 3, 3, activation='relu', )(c1)

    c2 = Convolution2D(64, 1, 1, activation='relu', border_mode='same')(m)
    c2 = Convolution2D(64, 7, 1, activation='relu', border_mode='same')(c2)
    c2 = Convolution2D(64, 1, 7, activation='relu', border_mode='same')(c2)
    c2 = Convolution2D(96, 3, 3, activation='relu', border_mode='valid')(c2)
    c2 = BatchNormalization(axis=channel_axis)(c2)

    m2 = merge([c1, c2], mode='concat', concat_axis=channel_axis)

    p1 = MaxPooling2D((3, 3), strides=(2, 2), )(m2)
    p2 = Convolution2D(192, 3, 3, activation='relu', subsample=(2, 2))(m2)

    m3 = merge([p1, p2], mode='concat', concat_axis=channel_axis)
    m3 = BatchNormalization(axis=channel_axis)(m3)
    m3 = Activation('relu')(m3)
    return m3


def inception_A(input):
    if K.image_dim_ordering() == "th":
        channel_axis = 1
    else:
        channel_axis = -1

    a1 = AveragePooling2D((3, 3), strides=(1, 1), border_mode='same')(input)
    a1 = Convolution2D(96, 1, 1, activation='relu', border_mode='same')(a1)

    a2 = Convolution2D(96, 1, 1, activation='relu', border_mode='same')(input)

    a3 = Convolution2D(64, 1, 1, activation='relu', border_mode='same')(input)
    a3 = Convolution2D(96, 3, 3, activation='relu', border_mode='same')(a3)

    a4 = Convolution2D(64, 1, 1, activation='relu', border_mode='same')(input)
    a4 = Convolution2D(96, 3, 3, activation='relu', border_mode='same')(a4)
    a4 = BatchNormalization(axis=channel_axis)(a4)
    a4 = Convolution2D(96, 3, 3, activation='relu', border_mode='same')(a4)

    m = merge([a1, a2, a3, a4], mode='concat', concat_axis=channel_axis)
    m = BatchNormalization(axis=channel_axis)(m)
    m = Activation('relu')(m)
    return m


def inception_B(input):
    if K.image_dim_ordering() == "th":
        channel_axis = 1
    else:
        channel_axis = -1

    b1 = AveragePooling2D((3, 3), strides=(1, 1), border_mode='same')(input)
    b1 = Convolution2D(128, 1, 1, activation='relu', border_mode='same')(b1)

    b2 = Convolution2D(384, 1, 1, activation='relu', border_mode='same')(input)

    b3 = Convolution2D(192, 1, 1, activation='relu', border_mode='same')(input)
    b3 = Convolution2D(224, 1, 7, activation='relu', border_mode='same')(b3)
    b3 = Convolution2D(256, 1, 7, activation='relu', border_mode='same')(b3)
    b3 = BatchNormalization(axis=channel_axis)(b3)

    b4 = Convolution2D(192, 1, 1, activation='relu', border_mode='same')(input)
    b4 = Convolution2D(192, 1, 7, activation='relu', border_mode='same')(b4)
    b4 = Convolution2D(224, 7, 1, activation='relu', border_mode='same')(b4)
    b4 = BatchNormalization(axis=channel_axis)(b4)
    b4 = Convolution2D(224, 1, 7, activation='relu', border_mode='same')(b4)
    b4 = Convolution2D(256, 7, 1, activation='relu', border_mode='same')(b4)

    m = merge([b1, b2, b3, b4], mode='concat', concat_axis=channel_axis)
    m = BatchNormalization(axis=channel_axis)(m)
    m = Activation('relu')(m)
    return m


def inception_C(input):
    if K.image_dim_ordering() == "th":
        channel_axis = 1
    else:
        channel_axis = -1

    c1 = AveragePooling2D((3, 3), strides=(1, 1), border_mode='same')(input)
    c1 = Convolution2D(256, 1, 1, activation='relu', border_mode='same')(c1)

    c2 = Convolution2D(256, 1, 1, activation='relu', border_mode='same')(input)

    c3 = Convolution2D(384, 1, 1, activation='relu', border_mode='same')(input)
    c3_1 = Convolution2D(256, 1, 3, activation='relu', border_mode='same')(c3)
    c3_2 = Convolution2D(256, 3, 1, activation='relu', border_mode='same')(c3)

    c4 = Convolution2D(384, 1, 1, activation='relu', border_mode='same')(input)
    c4 = Convolution2D(192, 1, 3, activation='relu', border_mode='same')(c4)
    c4 = Convolution2D(224, 3, 1, activation='relu', border_mode='same')(c4)
    c4 = BatchNormalization(axis=channel_axis)(c4)
    c4_1 = Convolution2D(256, 3, 1, activation='relu', border_mode='same')(c4)
    c4_2 = Convolution2D(256, 1, 3, activation='relu', border_mode='same')(c4)

    m = merge([c1, c2, c3_1, c3_2, c4_1, c4_2], mode='concat', concat_axis=channel_axis)
    m = BatchNormalization(axis=channel_axis)(m)
    m = Activation('relu')(m)
    return m


def reduction_A(input, k=192, l=224, m=256, n=384):
    if K.image_dim_ordering() == "th":
        channel_axis = 1
    else:
        channel_axis = -1

    r1 = MaxPooling2D((3, 3), strides=(2, 2))(input)

    r2 = Convolution2D(n, 3, 3, activation='relu', subsample=(2, 2))(input)

    r3 = Convolution2D(k, 1, 1, activation='relu', border_mode='same')(input)
    r3 = Convolution2D(l, 3, 3, activation='relu', border_mode='same')(r3)
    r3 = BatchNormalization(axis=channel_axis)(r3)
    r3 = Convolution2D(m, 3, 3, activation='relu', subsample=(2, 2))(r3)

    m = merge([r1, r2, r3], mode='concat', concat_axis=channel_axis)
    m = BatchNormalization(axis=channel_axis)(m)
    m = Activation('relu')(m)
    return m


def reduction_B(input):
    if K.image_dim_ordering() == "th":
        channel_axis = 1
    else:
        channel_axis = -1

    r1 = MaxPooling2D((3, 3), strides=(2, 2))(input)

    r2 = Convolution2D(192, 1, 1, activation='relu')(input)
    r2 = Convolution2D(192, 3, 3, activation='relu', subsample=(2, 2))(r2)

    r3 = Convolution2D(256, 1, 1, activation='relu', border_mode='same')(input)
    r3 = BatchNormalization(axis=channel_axis)(r3)
    r3 = Convolution2D(256, 1, 7, activation='relu', border_mode='same')(r3)
    r3 = Convolution2D(320, 7, 1, activation='relu', border_mode='same')(r3)
    r3 = BatchNormalization(axis=channel_axis)(r3)
    r3 = Convolution2D(320, 3, 3, activation='relu', border_mode='valid',
                       subsample=(2, 2))(r3)

    m = merge([r1, r2, r3], mode='concat', concat_axis=channel_axis)
    m = BatchNormalization(axis=channel_axis)(m)
    m = Activation('relu')(m)
    return m


def create_inception_v4(nb_classes=8):
    '''
    Creates a inception v4 network

    :param nb_classes: number of classes
    :return: Keras Model with 1 input and 1 output
    '''

    if K.image_dim_ordering() == 'th':
        init = Input((3, 299, 299))
    else:
        init = Input((299, 299, 3))

    # Input Shape is 299 x 299 x 3 (tf) or 3 x 299 x 299 (th)
    x = inception_stem(init)

    # 4 x Inception A
    for i in range(4):
        x = inception_A(x)

    # Reduction A
    x = reduction_A(x)

    # 7 x Inception B
    for i in range(7):
        x = inception_B(x)

    # Reduction B
    x = reduction_B(x)

    # 3 x Inception C
    for i in range(3):
        x = inception_C(x)

    # Average Pooling
    x = AveragePooling2D((8, 8))(x)

    # Dropout
    x = Dropout(0.8)(x)
    x = Flatten()(x)

    # Output
    out = Dense(output_dim=nb_classes, activation='softmax')(x)

    model = Model(init, out, name='Inception-v4')

    return model


if __name__ == '__main__':
    valX, valY, valX_id = process_validation((224, 224))
    trX, trY, trX_id = process_train((224, 224))
    tmp = np.zeros((300, 300, 3))
    new = np.zeros((299, 299, 3))
    valX_new = np.zeros((len(valX), 299, 299, 3))
    trX_new = np.zeros((len(trX), 299, 299, 3))
    #import pdb
    for img in range(len(valX)):
        tmp[:, :, 0] = np.pad(trX[img][:, :, 0], 38, 'constant')
        new[:, :, 0] = tmp[1:300, 1:300, 0]

        tmp[:, :, 1] = np.pad(trX[img][:, :, 1], 38, 'constant')
        new[:, :, 1] = tmp[1:300, 1:300, 1]

        tmp[:, :, 2] = np.pad(trX[img][:, :, 2], 38, 'constant')
        new[:, :, 2] = tmp[1:300, 1:300, 2]

        trX_new[img] = new

    del trX

    model = create_inception_v4()
    model.compile(optimizer='adam', loss='categorical_crossentropy',
                  metrics=['accuracy'])

    early_stopping = EarlyStopping(monitor='val_loss', patience=5, verbose=1)
    reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=3,
                                  min_lr=1e-10, mode='auto', verbose=1)

    model.fit(trX_new, trY, batch_size=8, nb_epoch=100, shuffle=True,
              validation_data=(valX_new, valY),
              callbacks=[early_stopping, reduce_lr], verbose=1)

    # Evaluation.
    scores = model.evaluate(valX_new, valY, verbose=0)
    print('Validation loss:', scores[0])
    print('Validation accuracy:', scores[1])

    teX, teX_id = process_test((224, 224))
    teX_new = np.zeros((len(teX), 299, 299, 3))
    for img in range(len(teX)):
        tmp[:, :, 0] = np.pad(teX[img][:, :, 0], 38, 'constant')
        new[:, :, 0] = tmp[1:300, 1:300, 0]

        tmp[:, :, 1] = np.pad(teX[img][:, :, 1], 38, 'constant')
        new[:, :, 1] = tmp[1:300, 1:300, 1]

        tmp[:, :, 2] = np.pad(teX[img][:, :, 2], 38, 'constant')
        new[:, :, 2] = tmp[1:300, 1:300, 2]

        teX_new[img] = new

    del teX

    y_pred = model.predict(teX_new)
    df = pd.DataFrame(y_pred,
                      columns=['ALB', 'BET', 'DOL', 'LAG',
                               'NoF', 'OTHER', 'SHARK', 'YFT'])
    df.insert(0, 'image', teX_id)
    df.to_csv('submission_inception_v4.csv', index=False)
    #test_scores = model.evaluate(teX_new, teY, verbose=0)
    #print("Test loss: {}".format(test_scores[0]))
    #print("Test accuracy: {}".format(test_scores[1]))
