"""
Face keypoints estimation model
"""
import tensorflow as tf

from functools import partial

from .estimation_utils import decode_keypoints_argmax_2d, decode_keypoints_compensated_2d


def Heatmaps():
    """Heatmap generating encoder
    """
    def heatmaps(x_in):
        x = inputs = tf.keras.layers.Input(x_in.shape[1:])
        x = tf.keras.layers.Conv2DTranspose(64, 4, strides=2, padding='same', use_bias=False, name='transpose_conv_1')(x)
        x = tf.keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.1, name='bn_1')(x)
        x = tf.keras.layers.ReLU(name='relu_1')(x)
        x = tf.keras.layers.Conv2DTranspose(64, 4, strides=2, padding='same', use_bias=False, name='transpose_conv_2')(x)
        x = tf.keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.1, name='bn_2')(x)
        x = tf.keras.layers.ReLU(name='relu_2')(x)
        x = tf.keras.layers.Conv2DTranspose(64, 4, strides=2, padding='same', use_bias=False, name='transpose_conv_3')(x)
        x = tf.keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.1, name='bn_3')(x)
        x = tf.keras.layers.ReLU(name='relu_3')(x)
        x = outputs = tf.keras.layers.Conv2D(98, 1, strides=1, name='conv_4')(x)
        return tf.keras.Model(inputs=inputs, outputs=outputs, name='heatmaps')(x_in)
    return heatmaps


def FaceKeypointEstimator(backbone, input_size, hm_dims=None, imagenet=False, training=False, decode='argmax', delta=4):
    """Fce keypoint estimation model

    Args:
        backbone (str): efficientnet type backbone, one of {efficientnet-b0, efficientnet-b1, efficientnet-b2}
        input_size (tuple): input image size in first two dimensions
        hm_dims (tuple, optional): generated heatmap dimensions in x and y axes. Defaults to None.
        imagenet (bool, optional): flag wether to load imagenet efficientnet weights. Defaults to False.
        training (bool, optional): training flag, if False - testing mode is enabled. Defaults to False.
        decode (str, optional): heatmap decoding method. If None - no decoding is performed and raw heatmaps are outputted,
                                otherwise one of {'argmax', 'compensated'}. Defaults to 'argmax'.
        delta (int, optional): delta value for compensated heatmap decoding. Defaults to 4.

    Returns:
        tf.keras.Model: constructed model
    """
    # backbones
    backbone_dict = {
        'efficientnet-b0': tf.keras.applications.efficientnet.EfficientNetB0,
        'efficientnet-b1': tf.keras.applications.efficientnet.EfficientNetB1,
        'efficientnet-b2': tf.keras.applications.efficientnet.EfficientNetB2,
    }
    x = inputs = tf.keras.layers.Input((*input_size, 3))
    x = backbone_dict[backbone](include_top=False, weights='imagenet' if imagenet else None)(x)
    x = tf.keras.layers.Dropout(0.2)(x)
    x = Heatmaps()(x)
    # if training is enabled or decoding is disabled, return raw heatmaps
    if not training and decode is not None:
        decode_method = {
            'argmax': decode_keypoints_argmax_2d,
            'compensated': partial(decode_keypoints_compensated_2d, delta=delta)
        }
        # decode heatmaps and norm coordinates to <0, 1> range
        x = tf.keras.layers.Lambda(lambda x: decode_method[decode](x))(x)
        x = tf.keras.layers.Lambda(lambda x: x / hm_dims)(x)
    return tf.keras.Model(inputs=inputs, outputs=x, name='FaceKeypointEstimator')


def load_model(weights_path, *args, **kwargs):
    """Load face estimation model (this serves mainly as a compatibility bridge between TF versions
       as models trained in TF 2.5 can't be automatically loaded with model.load_weights(...) in TF 2.8
       when some layers are left out).

    Args:
        weights_path (str): path to trained weights
        *args: model arguments
        **kwargs: model keyword arguments

    Returns:
        tf.keras.Model: Loaded face pose estimation model
    """
    # construct model in training mode and load its weights
    model = FaceKeypointEstimator(*args, **kwargs, training=True)
    model.load_weights(weights_path).expect_partial()
    # construct inference mode model
    inf_model = FaceKeypointEstimator(*args, **kwargs, training=False)
    # copy weights from models based on layer names
    for layer in model.layers:
        if len(layer.get_weights()):
            inf_model.get_layer(layer.name).set_weights(layer.get_weights())
    return inf_model
