import sys, os
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from utils import load_externals

from mobilenets import *
from mobilenets import _obtain_input_shape, __conv_block, __depthwise_conv_block

from keras.layers import Lambda

def scaling_tf(X, input_range_type):
    """
    Convert to [-1, 1].
    """
    if input_range_type == 1:
        # The input data range is [0, 1]. Convert to [-1, 1] by
        X = X - 0.5
        X = X * 2.
    elif input_range_type == 2:
        # The input data range is [-0.5, 0.5]. Convert to [-1,1] by
        X = X * 2.
    elif input_range_type == 3:
        # The input data range is [-1, 1].
        X = X

    return X


def __create_mobilenet(classes, img_input, include_top, alpha, depth_multiplier, dropout, pooling, logits):
    ''' Creates a MobileNet model with specified parameters
    Args:
        classes: Number of output classes
        img_input: Input tensor or layer
        include_top: Flag to include the last dense layer
        alpha: width multiplier of the MobileNet.
        depth_multiplier: depth multiplier for depthwise convolution
                          (also called the resolution multiplier)
        dropout: dropout rate
        pooling: Optional pooling mode for feature extraction
            when `include_top` is `False`.
            - `None` means that the output of the model will be
                the 4D tensor output of the
                last convolutional layer.
            - `avg` means that global average pooling
                will be applied to the output of the
                last convolutional layer, and thus
                the output of the model will be a 2D tensor.
            - `max` means that global max pooling will
                be applied.
    Returns: a Keras Model
    '''

    x = __conv_block(img_input, 32, alpha, strides=(2, 2))
    x = __depthwise_conv_block(x, 64, alpha, depth_multiplier, id=1)

    x = __depthwise_conv_block(x, 128, alpha, depth_multiplier, strides=(2, 2), id=2)
    x = __depthwise_conv_block(x, 128, alpha, depth_multiplier, id=3)

    x = __depthwise_conv_block(x, 256, alpha, depth_multiplier, strides=(2, 2), id=4)
    x = __depthwise_conv_block(x, 256, alpha, depth_multiplier, id=5)

    x = __depthwise_conv_block(x, 512, alpha, depth_multiplier, strides=(2, 2), id=6)
    x = __depthwise_conv_block(x, 512, alpha, depth_multiplier, id=7)
    x = __depthwise_conv_block(x, 512, alpha, depth_multiplier, id=8)
    x = __depthwise_conv_block(x, 512, alpha, depth_multiplier, id=9)
    x = __depthwise_conv_block(x, 512, alpha, depth_multiplier, id=10)
    x = __depthwise_conv_block(x, 512, alpha, depth_multiplier, id=11)

    x = __depthwise_conv_block(x, 1024, alpha, depth_multiplier, strides=(2, 2), id=12)
    x = __depthwise_conv_block(x, 1024, alpha, depth_multiplier, id=13)

    if include_top:
        if K.image_data_format() == 'channels_first':
            shape = (int(1024 * alpha), 1, 1)
        else:
            shape = (1, 1, int(1024 * alpha))

        x = GlobalAveragePooling2D()(x)
        x = Reshape(shape, name='reshape_1')(x)
        x = Dropout(dropout, name='dropout')(x)
        x = Convolution2D(classes, (1, 1), padding='same', name='conv_preds')(x)
        # Reshape from (?, 1, 1, 1000)  to (?, 1000)
        x = Reshape((classes,), name='reshape_2')(x)

        # Move Reshape before Actionvation. Otherwise, Cleverhans gets confused in fetching logits output.
        if not logits:
            x = Activation('softmax', name='activation')(x)
    else:
        if pooling == 'avg':
            x = GlobalAveragePooling2D()(x)
        elif pooling == 'max':
            x = GlobalMaxPooling2D()(x)
    return x


def MobileNets(input_shape=None, alpha=1.0, depth_multiplier=1,
               dropout=1e-3, include_top=True, weights='imagenet',
               input_tensor=None, pooling=None, classes=1000,
               logits=False, input_range_type=1, pre_filter=lambda x:x):
    ''' Instantiate the MobileNet architecture.
        Note that only TensorFlow is supported for now,
        therefore it only works with the data format
        `image_data_format='channels_last'` in your Keras config
        at ~/.keras/keras.json.
        # Arguments
            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 inputs channels,
                and width and height should be no smaller than 32.
                E.g. `(200, 200, 3)` would be one valid value.
            alpha: width multiplier of the MobileNet.
            depth_multiplier: depth multiplier for depthwise convolution
                (also called the resolution multiplier)
            dropout: dropout rate
            include_top: whether to include the fully-connected
                layer at the top of the network.
            weights: `None` (random initialization) or
                `imagenet` (ImageNet weights)
            input_tensor: optional Keras tensor (i.e. output of `layers.Input()`)
                to use as image input for the model.
            pooling: Optional pooling mode for feature extraction
                when `include_top` is `False`.
                - `None` means that the output of the model will be
                    the 4D tensor output of the
                    last convolutional layer.
                - `avg` means that global average pooling
                    will be applied to the output of the
                    last convolutional layer, and thus
                    the output of the model will be a 2D tensor.