import tensorflow as tf
from tensorflow.keras import Model
from tensorflow.keras.layers import (
    Dense,
    Dropout,
    Flatten,
    Input,
)
from tensorflow.keras.applications import (
    MobileNetV2,
    ResNet50
)
from .layers import (
    BatchNormalization,
    ArcMarginPenaltyLogists
)


def _regularizer(weights_decay=5e-4):
    return tf.keras.regularizers.l2(weights_decay)


def Backbone(backbone_type='ResNet50', use_pretrain=True):
    """Backbone Model"""
    weights = None
    if use_pretrain:
        weights = 'imagenet'

    def backbone(x_in):
        if backbone_type == 'ResNet50':
            return ResNet50(input_shape=x_in.shape[1:], include_top=False,
                            weights=weights)(x_in)
        elif backbone_type == 'MobileNetV2':
            return MobileNetV2(input_shape=x_in.shape[1:], include_top=False,
                               weights=weights)(x_in)
        else:
            raise TypeError('backbone_type error!')
    return backbone


def OutputLayer(embd_shape, w_decay=5e-4, name='OutputLayer'):
    """Output Later"""
    def output_layer(x_in):
        x = inputs = Input(x_in.shape[1:])
        x = BatchNormalization()(x)
        x = Dropout(rate=0.5)(x)
        x = Flatten()(x)
        x = Dense(embd_shape, kernel_regularizer=_regularizer(w_decay))(x)
        x = BatchNormalization()(x)
        return Model(inputs, x, name=name)(x_in)
    return output_layer


def ArcHead(num_classes, margin=0.5, logist_scale=64, name='ArcHead'):
    """Arc Head"""
    def arc_head(x_in, y_in):
        x = inputs1 = Input(x_in.shape[1:])
        y = Input(y_in.shape[1:])
        x = ArcMarginPenaltyLogists(num_classes=num_classes,
                                    margin=margin,
                                    logist_scale=logist_scale)(x, y)
        return Model((inputs1, y), x, name=name)((x_in, y_in))
    return arc_head


def NormHead(num_classes, w_decay=5e-4, name='NormHead'):
    """Norm Head"""
    def norm_head(x_in):
        x = inputs = Input(x_in.shape[1:])
        x = Dense(num_classes, kernel_regularizer=_regularizer(w_decay))(x)
        return Model(inputs, x, name=name)(x_in)
    return norm_head


def ArcFaceModel(size=None, channels=3, num_classes=None, name='arcface_model',
                 margin=0.5, logist_scale=64, embd_shape=512,
                 head_type='ArcHead', backbone_type='ResNet50',
                 w_decay=5e-4, use_pretrain=True, training=False):
    """Arc Face Model"""
    x = inputs = Input([size, size, channels], name='input_image')

    x = Backbone(backbone_type=backbone_type, use_pretrain=use_pretrain)(x)

    embds = OutputLayer(embd_shape, w_decay=w_decay)(x)
    
    embds = embds / tf.norm(embds, axis=1, keepdims=True)

    if training:
        assert num_classes is not None
        labels = Input([], name='label')
        if head_type == 'ArcHead':
            logist = ArcHead(num_classes=num_classes, margin=margin,
                             logist_scale=logist_scale)(embds, labels)
        else:
            logist = NormHead(num_classes=num_classes, w_decay=w_decay)(embds)
        return Model((inputs, labels), logist, name=name)
    else:
        return Model(inputs, embds, name=name)

tensorflow.norm tensorflow.keras.regularizers.l2 tensorflow.keras.Model tensorflow.keras.applications.ResNet50 tensorflow.keras.layers.Dropout tensorflow.keras.layers.Flatten tensorflow.keras.applications.MobileNetV2 tensorflow.keras.layers.Input
[(20, 'tensorflow.keras.regularizers.l2', 'tf.keras.regularizers.l2', (['weights_decay'], {}), True, 'import tensorflow as tf\n'), (80, 'tensorflow.keras.layers.Input', 'Input', (['[size, size, channels]'], {'name': '"""input_image"""'}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten, Input\n'), (44, 'tensorflow.keras.layers.Input', 'Input', (['x_in.shape[1:]'], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten, Input\n'), (57, 'tensorflow.keras.layers.Input', 'Input', (['x_in.shape[1:]'], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten, Input\n'), (58, 'tensorflow.keras.layers.Input', 'Input', (['y_in.shape[1:]'], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten, Input\n'), (69, 'tensorflow.keras.layers.Input', 'Input', (['x_in.shape[1:]'], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten, Input\n'), (86, 'tensorflow.norm', 'tf.norm', (['embds'], {'axis': '(1)', 'keepdims': '(True)'}), True, 'import tensorflow as tf\n'), (90, 'tensorflow.keras.layers.Input', 'Input', (['[]'], {'name': '"""label"""'}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten, Input\n'), (96, 'tensorflow.keras.Model', 'Model', (['(inputs, labels)', 'logist'], {'name': 'name'}), False, 'from tensorflow.keras import Model\n'), (98, 'tensorflow.keras.Model', 'Model', (['inputs', 'embds'], {'name': 'name'}), False, 'from tensorflow.keras import Model\n'), (46, 'tensorflow.keras.layers.Dropout', 'Dropout', ([], {'rate': '(0.5)'}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten, Input\n'), (47, 'tensorflow.keras.layers.Flatten', 'Flatten', ([], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten, Input\n'), (50, 'tensorflow.keras.Model', 'Model', (['inputs', 'x'], {'name': 'name'}), False, 'from tensorflow.keras import Model\n'), (62, 'tensorflow.keras.Model', 'Model', (['(inputs1, y)', 'x'], {'name': 'name'}), False, 'from tensorflow.keras import Model\n'), (71, 'tensorflow.keras.Model', 'Model', (['inputs', 'x'], {'name': 'name'}), False, 'from tensorflow.keras import Model\n'), (31, 'tensorflow.keras.applications.ResNet50', 'ResNet50', ([], {'input_shape': 'x_in.shape[1:]', 'include_top': '(False)', 'weights': 'weights'}), False, 'from tensorflow.keras.applications import MobileNetV2, ResNet50\n'), (34, 'tensorflow.keras.applications.MobileNetV2', 'MobileNetV2', ([], {'input_shape': 'x_in.shape[1:]', 'include_top': '(False)', 'weights': 'weights'}), False, 'from tensorflow.keras.applications import MobileNetV2, ResNet50\n')]