# D:/workplace/python
# -*- coding: utf-8 -*-
# @File  :temp_mobilenetv3.py
# @Author:Guido    LuXiaohao
# @Date  :2020/4/9
# @Software:PyCharm


from keras.models import Model
from keras.layers import Input, Conv2D, GlobalAveragePooling2D, Reshape, Dropout, Dense, MaxPooling2D, concatenate,AveragePooling2D
from keras.utils.vis_utils import plot_model
from nets.mobilenet.mobilenet_base import MobileNetBase
from keras import backend as K
from keras.engine.topology import Layer
from keras.layers import activations, initializers, regularizers, constraints, Lambda
from keras.engine import InputSpec
import tensorflow as tf


class AMSoftmax(Layer):
    def __init__(self, units, s, m,
                 kernel_initializer='glorot_uniform',
                 kernel_regularizer=None,
                 kernel_constraint=None,
                 **kwargs
                 ):
        if 'input_shape' not in kwargs and 'input_dim' in kwargs:
            kwargs['input_shape'] = (kwargs.pop('input_dim'),)
        super(AMSoftmax, self).__init__(**kwargs)
        self.units = units
        self.s = s
        self.m = m
        self.kernel_initializer = initializers.get(kernel_initializer)
        self.kernel_regularizer = regularizers.get(kernel_regularizer)
        self.kernel_constraint = constraints.get(kernel_constraint)
        self.input_spec = InputSpec(min_ndim=2)
        self.supports_masking = True


    def build(self, input_shape):
        assert len(input_shape) >= 2
        input_dim = input_shape[-1]

        self.kernel = self.add_weight(shape=(input_dim, self.units),
                                      initializer=self.kernel_initializer,
                                      name='kernel',
                                      regularizer=self.kernel_regularizer,
                                      constraint=self.kernel_constraint)
        self.bias = None

        self.input_spec = InputSpec(min_ndim=2, axes={-1: input_dim})
        self.built = True


    def call(self, inputs, **kwargs):
        inputs = tf.nn.l2_normalize(inputs, dim=-1)
        self.kernel = tf.nn.l2_normalize(self.kernel, dim=(0, 1))   # W归一化

        dis_cosin = K.dot(inputs, self.kernel)
        psi = dis_cosin - self.m

        e_costheta = K.exp(self.s * dis_cosin)
        e_psi = K.exp(self.s * psi)
        sum_x = K.sum(e_costheta, axis=-1, keepdims=True)

        temp = e_psi - e_costheta
        temp = temp + sum_x

        output = e_psi / temp
        return output


def amsoftmax_loss(y_true, y_pred):
    d1 = K.sum(y_true * y_pred, axis=-1)
    d1 = K.log(K.clip(d1, K.epsilon(), None))
    loss = -K.mean(d1, axis=-1)
    return loss


class MobileNetV3_Small(MobileNetBase):
    def __init__(self, shape, n_class, alpha=1.0, include_top=True):
        """Init.

        # Arguments
            input_shape: An integer or tuple/list of 3 integers, shape
                of input tensor.
            n_class: Integer, number of classes.
            alpha: Integer, width multiplier.
            include_top: if inculde classification layer.

        # Returns
            MobileNetv3 model.
        """
        super(MobileNetV3_Small, self).__init__(shape, n_class, alpha)
        self.include_top = include_top

    def build(self, plot=False):
        """build MobileNetV3 Small.

        # Arguments
            plot: Boolean, weather to plot model.

        # Returns
            model: Model, model.
        """
        inputs = Input(shape=self.shape)

        x = self._conv_block(inputs, 16, (3, 3), strides=(2, 2), nl=None)

        x = self._bottleneck(x, 16, (3, 3), e=16, s=2, squeeze=True, nl=None)
        x = self._bottleneck(x, 24, (3, 3), e=72, s=2, squeeze=False, nl=None)
        x = self._bottleneck(x, 24, (3, 3), e=88, s=1, squeeze=False, nl=None)
        x = self._bottleneck(x, 40, (5, 5), e=96, s=2, squeeze=True, nl=None)
        x = self._bottleneck(x, 40, (5, 5), e=240, s=1, squeeze=True, nl=None)
        x = self._bottleneck(x, 40, (5, 5), e=240, s=1, squeeze=True, nl=None)
        x = self._bottleneck(x, 48, (5, 5), e=120, s=1, squeeze=True, nl=None)
        x = self._bottleneck(x, 48, (5, 5), e=144, s=1, squeeze=True, nl=None)
        x = self._bottleneck(x, 96, (5, 5), e=288, s=2, squeeze=True, nl=None)
        x = self._bottleneck(x, 96, (5, 5), e=576, s=1, squeeze=True, nl=None)
        x = self._bottleneck(x, 96, (5, 5), e=576, s=1, squeeze=True, nl=None)

        x = self._conv_block(x, 576, (1, 1), strides=(1, 1), nl=None)
        x = GlobalAveragePooling2D()(x)
        x = Reshape((1, 1, 576))(x)
        x = Conv2D(1280, (1, 1), padding='same')(x)
        x = self._return_activation(x, None)
        # x = Dropout(0.5)(x)

        if self.include_top:
            x = Conv2D(self.n_class, (1, 1), padding='same', activation='softmax')(x)
            x = Reshape((self.n_class,), name='reshape_11')(x)

        model = Model(inputs, x)
        model.summary()
        if plot:
            plot_model(model, to_file='images/MobileNetv3_small.png', show_shapes=True)

        return model

# A = MobileNetV3_Small((256, 256, 3), 4)
# model = A.build()
# model.save(r"C:\Users\Administrator\Desktop\1.h5")