# -*- coding: utf-8 -*-
"""Human Rights Archive (HRA) VGG16-Places365 model for Keras

◾◾◾◾ Fine-tuning acc. ◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾
◾                                                     ◾
◾ VGG16-Places365_avg  =>  30.00%                     ◾
◾                                                     ◾
◾ VGG16-Places365_flatten  =>  28.51%                 ◾
◾                                                     ◾
◾ VGG16-Places365_max  =>  26.66%                     ◾
◾                                                     ◾
◾◾◾◾ Fine-tuning coverage ◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾
◾                                                     ◾
◾ VGG16-Places365_avg  =>  32%                        ◾
◾                                                     ◾
◾ VGG16-Places365_flatten  =>  52%                    ◾
◾                                                     ◾
◾ VGG16-Places365_max  =>  51%                        ◾
◾                                                     ◾
◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾◾

"""

from __future__ import division, print_function
import os

import warnings
import numpy as np

from keras import backend as K
from keras.utils.data_utils import get_file
from keras.layers import Input
from keras.layers.core import Dense, Flatten
from keras.models import Model
from keras.layers.core import Dropout
from keras.layers import GlobalAveragePooling2D
from keras.layers import GlobalMaxPooling2D
from applications.vgg16_places_365 import VGG16_Places365
from keras.applications.imagenet_utils import _obtain_input_shape
from keras.engine.topology import get_source_inputs
from keras.optimizers import SGD

# ======================  ------------------------  ======================  ------------------------  ======================  ------------------------  ======================  ------------------------  ======================  ------------------------
# Weights obtained by re-training the last Dense layer
FEATURE_EXTRACTION_AVG_POOL_WEIGHTS_PATH = 'https://github.com/GKalliatakis/crispy-enigma/releases/download/0.1/feature_extraction_VGG16_Places365_avg_pool_weights_tf_dim_ordering_tf_kernels.h5'
FEATURE_EXTRACTION_AVG_POOL_fname = 'feature_extraction_VGG16_Places365_avg_pool_weights_tf_dim_ordering_tf_kernels'
FEATURE_EXTRACTION_FLATTEN_POOL_WEIGHTS_PATH = 'https://github.com/GKalliatakis/crispy-enigma/releases/download/0.1/feature_extraction_VGG16_Places365_flatten_pool_weights_tf_dim_ordering_tf_kernels.h5'
FEATURE_EXTRACTION_FLATTEN_POOL_fname = 'feature_extraction_VGG16_Places365_flatten_pool_weights_tf_dim_ordering_tf_kernels'
FEATURE_EXTRACTION_MAX_POOL_WEIGHTS_PATH = 'https://github.com/GKalliatakis/crispy-enigma/releases/download/0.1/feature_extraction_VGG16_Places365_max_pool_weights_tf_dim_ordering_tf_kernels.h5'
FEATURE_EXTRACTION_MAX_POOL_fname = 'feature_extraction_VGG16_Places365_max_pool_weights_tf_dim_ordering_tf_kernels'

AUGM_FEATURE_EXTRACTION_AVG_POOL_WEIGHTS_PATH = 'https://github.com/GKalliatakis/crispy-enigma/releases/download/0.1/augm_feature_extraction_VGG16_Places365_avg_pool_weights_tf_dim_ordering_tf_kernels.h5'
AUGM_FEATURE_EXTRACTION_AVG_POOL_fname = 'augm_feature_extraction_VGG16_Places365_avg_pool_weights_tf_dim_ordering_tf_kernels'
AUGM_FEATURE_EXTRACTION_FLATTEN_POOL_WEIGHTS_PATH = 'https://github.com/GKalliatakis/crispy-enigma/releases/download/0.1/augm_feature_extraction_VGG16_Places365_flatten_pool_weights_tf_dim_ordering_tf_kernels.h5'
AUGM_FEATURE_EXTRACTION_FLATTEN_POOL_fname = 'augm_feature_extraction_VGG16_Places365_flatten_pool_weights_tf_dim_ordering_tf_kernels'
AUGM_FEATURE_EXTRACTION_MAX_POOL_WEIGHTS_PATH = 'https://github.com/GKalliatakis/crispy-enigma/releases/download/0.1/augm_feature_extraction_VGG16_Places365_max_pool_weights_tf_dim_ordering_tf_kernels.h5'
AUGM_FEATURE_EXTRACTION_MAX_POOL_fname = 'augm_feature_extraction_VGG16_Places365_max_pool_weights_tf_dim_ordering_tf_kernels'

# ======================  ------------------------  ======================  ------------------------  ======================  ------------------------  ======================  ------------------------  ======================  ------------------------
# Weights obtained by un-freezing the two lower convolutional layers and retraining them
FINE_TUNING_AVG_POOL_WEIGHTS_PATH = 'https://github.com/GKalliatakis/crispy-enigma/releases/download/0.2/fine_tuning_VGG16_Places365_avg_pool_weights_tf_dim_ordering_tf_kernels.h5'
FINE_TUNING_AVG_POOL_fname = 'fine_tuning_VGG16_Places365_avg_pool_weights_tf_dim_ordering_tf_kernels'
FINE_TUNING_FLATTEN_POOL_WEIGHTS_PATH = 'https://github.com/GKalliatakis/crispy-enigma/releases/download/0.2/fine_tuning_VGG16_Places365_flatten_pool_weights_tf_dim_ordering_tf_kernels.h5'
FINE_TUNING_FLATTEN_POOL_fname = 'fine_tuning_VGG16_Places365_flatten_pool_weights_tf_dim_ordering_tf_kernels'
FINE_TUNING_MAX_POOL_WEIGHTS_PATH = 'https://github.com/GKalliatakis/crispy-enigma/releases/download/0.2/fine_tuning_VGG16_Places365_max_pool_weights_tf_dim_ordering_tf_kernels.h5'
FINE_TUNING_MAX_POOL_fname = 'fine_tuning_VGG16_Places365_max_pool_weights_tf_dim_ordering_tf_kernels'

AUGM_FINE_TUNING_AVG_POOL_WEIGHTS_PATH = 'https://github.com/GKalliatakis/crispy-enigma/releases/download/0.2/augm_fine_tuning_VGG16_Places365_avg_pool_weights_tf_dim_ordering_tf_kernels.h5'
AUGM_FINE_TUNING_AVG_POOL_fname = 'augm_fine_tuning_VGG16_Places365_avg_pool_weights_tf_dim_ordering_tf_kernels'
AUGM_FINE_TUNING_FLATTEN_POOL_WEIGHTS_PATH = 'https://github.com/GKalliatakis/crispy-enigma/releases/download/0.2/augm_fine_tuning_VGG16_Places365_flatten_pool_weights_tf_dim_ordering_tf_kernels.h5'
AUGM_FINE_TUNING_FLATTEN_POOL_fname = 'augm_fine_tuning_VGG16_Places365_flatten_pool_weights_tf_dim_ordering_tf_kernels'
AUGM_FINE_TUNING_MAX_POOL_WEIGHTS_PATH = 'https://github.com/GKalliatakis/crispy-enigma/releases/download/0.2/augm_fine_tuning_VGG16_Places365_max_pool_weights_tf_dim_ordering_tf_kernels.h5'
AUGM_FINE_TUNING_MAX_POOL_fname = 'augm_fine_tuning_VGG16_Places365_max_pool_weights_tf_dim_ordering_tf_kernels'

# ======================  ------------------------  ======================  ------------------------  ======================  ------------------------  ======================  ------------------------  ======================  ------------------------
# Weights obtained from removing the classification block from the fine-tuned models
FINE_TUNING_WEIGHTS_PATH_NO_TOP = 'https://github.com/GKalliatakis/crispy-enigma/releases/download/0.2/fine_tuning_VGG16_Places365_weights_tf_dim_ordering_tf_kernels_notop.h5'
FINE_TUNING_WEIGHTS_PATH_NO_TOP_fname ='fine_tuning_VGG16_Places365_weights_tf_dim_ordering_tf_kernels_notop'
# ======================  ------------------------  ======================  ------------------------  ======================  ------------------------  ======================  ------------------------  ======================  ------------------------

def HRA_VGG16_Places365(include_top=True, weights='HRA',
                        input_tensor=None, input_shape=None,
                        mode= 'fine_tuning',
                        pooling_mode='avg',
                        classes=9,
                        data_augm_enabled=False):
    """Instantiates the VGG16_Places365 architecture fine-tuned (2 steps) on Human Rights Archive dataset.

    Optionally loads weights pre-trained on Human Rights Archive Database.

    # Arguments
        include_top: whether to include the 3 fully-connected
            layers at the top of the network.
        weights: one of `None` (random initialization),
            'HRA' (pre-training on Human Rights Archive),
            or the path to the weights file to be loaded.
        input_tensor: optional Keras tensor (i.e. output of `layers.Input()`)
            to use as image input for the model.
        input_shape: optional shape tuple, only to be specified
            if `include_top` is False (otherwise the input shape
            has to be `(224, 224, 3)` (with `channels_last` data format)
            or `(3, 224, 224)` (with `channels_first` data format).
            It should have exactly 3 input channels,
            and width and height should be no smaller than 48.
            E.g. `(200, 200, 3)` would be one valid value.
        mode: one of `feature_extraction` (freeze all but the penultimate layer and re-train the last Dense layer)
            or `fine_tuning` (unfreeze the lower convolutional layers and retrain more layers).
        pooling_mode: Pooling mode that will be applied to the output of the last convolutional layer of the original model
            and thus the output of the model will be a 2D tensor.
            - `avg` means that global average pooling_mode operation for spatial data will be applied.
            - `max` means that global max pooling_mode operation for spatial data will be applied.
            - `flatten` means that the output of the the last convolutional
                layer of the original model will be flatten,
                resulting in a larger Dense layer afterwards.
        classes: optional number of classes to classify images into.
        data_augm_enabled: whether to use the augmented samples during training.

    # Returns
        A Keras model instance.

    # Raises
        ValueError: in case of invalid argument for `weights`, or invalid input shape
        """
    if not (weights in {'HRA', None} or os.path.exists(weights)):
        raise ValueError('The `weights` argument should be either '
                         '`None` (random initialization), `HRA` '
                         '(pre-training on Human Rights Archive), '
                         'or the path to the weights file to be loaded.')

    if not (mode in {'feature_extraction', 'fine_tuning'}):
        raise ValueError('The `mode` argument should be either '
                         '`feature_extraction` (freeze all but the penultimate layer and re-train the last Dense layer),'
                         'or `fine_tuning` (unfreeze the lower convolutional layers and retrain more layers). ')

    if not (pooling_mode in {'avg', 'max', 'flatten'}):
        raise ValueError('The `pooling_mode` argument should be either '
                         '`avg` (global average pooling_mode), `max` (global max pooling_mode)'
                         'or `flatten` (the output will be flatten). ')

    if mode == 'feature_extraction' and include_top is False:
        raise ValueError('The `include_top` argument can be set as false only '
                         'when the `mode` argument is `fine_tuning`. '
                         'If not, the returned model would have been literally the default '
                         'keras-applications model and not the one trained on HRA.')

    cache_subdir = 'HRA_models'

    # Determine proper input shape
    input_shape = _obtain_input_shape(input_shape,
                                      default_size=224,
                                      min_size=48,
                                      data_format=K.image_data_format(),
                                      require_flatten=include_top,
                                      weights=weights)

    if input_tensor is None:
        img_input = Input(shape=input_shape)
    else:
        if not K.is_keras_tensor(input_tensor):
            img_input = Input(tensor=input_tensor, shape=input_shape)
        else:
            img_input = input_tensor

    # Ensure that the model takes into account any potential predecessors of `input_tensor`.
    if input_tensor is not None:
        inputs = get_source_inputs(input_tensor)
    else:
        inputs = img_input


    # create the base pre-trained model
    base_model = VGG16_Places365(weights='places', include_top=False, input_tensor=img_input)
    x = base_model.output

    # Classification block - build a classifier model to put on top of the convolutional model
    if include_top:

        # add a global spatial pooling_mode layer or flatten the obtained output from the original model
        if pooling_mode == 'avg':
            x = GlobalAveragePooling2D(name='GAP')(x)
        elif pooling_mode == 'max':
            x = GlobalMaxPooling2D(name='GMP')(x)
        elif pooling_mode == 'flatten':
            x = Flatten(name='FLATTEN')(x)

        # add a fully-connected layer
        x = Dense(256, activation='relu', name='FC1')(x)

        # When random init is enabled, we want to include Dropout,
        # otherwise when loading a pre-trained HRA model we want to omit that layer,
        # so the visualisations are done properly (there is an issue if it is included)
        if weights is None:
            x = Dropout(0.5,name='DROPOUT')(x)
        # and a logistic layer with the number of classes defined by the `classes` argument
        x = Dense(classes, activation='softmax', name='PREDICTIONS')(x)

    model = Model(inputs=inputs, outputs=x, name='HRA-VGG16_Places365')


    # load weights
    if weights == 'HRA':
        if include_top:
            if mode == 'feature_extraction':
                for layer in base_model.layers:
                    layer.trainable = False

                model.compile(optimizer=SGD(lr=0.0001, momentum=0.9),
                              loss='categorical_crossentropy')

                if data_augm_enabled:
                    if pooling_mode == 'avg':
                        weights_path = get_file(AUGM_FEATURE_EXTRACTION_AVG_POOL_fname,
                                                AUGM_FEATURE_EXTRACTION_AVG_POOL_WEIGHTS_PATH,
                                                cache_subdir=cache_subdir)
                    elif pooling_mode == 'flatten':
                        weights_path = get_file(AUGM_FEATURE_EXTRACTION_FLATTEN_POOL_fname,
                                                AUGM_FEATURE_EXTRACTION_FLATTEN_POOL_WEIGHTS_PATH,
                                                cache_subdir=cache_subdir)

                    elif pooling_mode == 'max':
                        weights_path = get_file(AUGM_FEATURE_EXTRACTION_MAX_POOL_fname,
                                                AUGM_FEATURE_EXTRACTION_MAX_POOL_WEIGHTS_PATH,
                                                cache_subdir=cache_subdir)

                else:
                    if pooling_mode == 'avg':
                        weights_path = get_file(FEATURE_EXTRACTION_AVG_POOL_fname,
                                                FEATURE_EXTRACTION_AVG_POOL_WEIGHTS_PATH,
                                                cache_subdir=cache_subdir)
                    elif pooling_mode == 'flatten':
                        weights_path = get_file(FEATURE_EXTRACTION_FLATTEN_POOL_fname,
                                                FEATURE_EXTRACTION_FLATTEN_POOL_WEIGHTS_PATH,
                                                cache_subdir=cache_subdir)


                    elif pooling_mode == 'max':
                        weights_path = get_file(FEATURE_EXTRACTION_MAX_POOL_fname,
                                                FEATURE_EXTRACTION_MAX_POOL_WEIGHTS_PATH,
                                                cache_subdir=cache_subdir)


            elif mode == 'fine_tuning':
                if pooling_mode == 'flatten':
                    for layer in model.layers[:17]:
                        layer.trainable = False
                    for layer in model.layers[17:]:
                        layer.trainable = True
                else:
                    for layer in model.layers[:16]:
                        layer.trainable = False
                    for layer in model.layers[16:]:
                        layer.trainable = True

                model.compile(optimizer=SGD(lr=0.0001, momentum=0.9),
                              loss='categorical_crossentropy',
                              metrics=['accuracy'])

                if data_augm_enabled:
                    if pooling_mode == 'avg':
                        weights_path = get_file(AUGM_FINE_TUNING_AVG_POOL_fname,
                                                AUGM_FINE_TUNING_AVG_POOL_WEIGHTS_PATH,
                                                cache_subdir=cache_subdir)
                    elif pooling_mode == 'flatten':
                        weights_path = get_file(AUGM_FINE_TUNING_FLATTEN_POOL_fname,
                                                AUGM_FINE_TUNING_FLATTEN_POOL_WEIGHTS_PATH,
                                                cache_subdir=cache_subdir)

                    elif pooling_mode == 'max':
                        weights_path = get_file(AUGM_FINE_TUNING_MAX_POOL_fname,
                                                AUGM_FINE_TUNING_MAX_POOL_WEIGHTS_PATH,
                                                cache_subdir=cache_subdir)

                else:
                    if pooling_mode == 'avg':
                        weights_path = get_file(FINE_TUNING_AVG_POOL_fname,
                                                FINE_TUNING_AVG_POOL_WEIGHTS_PATH,
                                                cache_subdir=cache_subdir)
                    elif pooling_mode == 'flatten':
                        weights_path = get_file(FINE_TUNING_FLATTEN_POOL_fname,
                                                FINE_TUNING_FLATTEN_POOL_WEIGHTS_PATH,
                                                cache_subdir=cache_subdir)


                    elif pooling_mode == 'max':
                        weights_path = get_file(FINE_TUNING_MAX_POOL_fname,
                                                FINE_TUNING_MAX_POOL_WEIGHTS_PATH,
                                                cache_subdir=cache_subdir)

        else:
            weights_path = get_file(FINE_TUNING_WEIGHTS_PATH_NO_TOP_fname,
                                    FINE_TUNING_WEIGHTS_PATH_NO_TOP,
                                    cache_subdir=cache_subdir)

        model.load_weights(weights_path)

    elif weights is not None:
        model.load_weights(weights)


    return model

