import tensorflow as tf

from functools import partial

from .estimation_utils import decode_keypoints_argmax_2d, decode_keypoints_compensated_2d


def SingleHandHeatmaps():
    """Single hand heatmap generation branch
    """
    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)(x)
        x = tf.keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.1)(x)
        x = tf.keras.layers.ReLU()(x)
        x = tf.keras.layers.Conv2DTranspose(64, 4, strides=2, padding='same', use_bias=False)(x)
        x = tf.keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.1)(x)
        x = tf.keras.layers.ReLU()(x)
        x = tf.keras.layers.Dropout(0.2)(x)
        x = tf.keras.layers.Conv2DTranspose(64, 4, strides=2, padding='same', use_bias=False)(x)
        x = tf.keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.1)(x)
        x = tf.keras.layers.ReLU()(x)
        x = outputs = tf.keras.layers.Conv2D(21, 1, strides=1)(x)
        return tf.keras.Model(inputs=inputs, outputs=outputs, name='SingleHandHeatmaps')(x_in)
    return heatmaps


def InteractingHeatmaps():
    """Interacting hands heatmap generation branch
    """
    def heatmaps(x_in):
        x = inputs = tf.keras.layers.Input(x_in.shape[1:])
        x = tf.keras.layers.Conv2DTranspose(128, 4, strides=2, padding='same', use_bias=False)(x)
        x = tf.keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.1)(x)
        x = tf.keras.layers.ReLU()(x)
        x = tf.keras.layers.Conv2DTranspose(128, 4, strides=2, padding='same', use_bias=False)(x)
        x = tf.keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.1, name='bn_2')(x)
        x = tf.keras.layers.ReLU()(x)
        x = tf.keras.layers.Dropout(0.2)(x)
        x = tf.keras.layers.Conv2DTranspose(128, 4, strides=2, padding='same', use_bias=False)(x)
        x = tf.keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.1, name='bn_3')(x)
        x = tf.keras.layers.ReLU()(x)
        x = outputs = tf.keras.layers.Conv2D(42, 1, strides=1)(x)
        return tf.keras.Model(inputs=inputs, outputs=outputs, name='InteractingHeatmaps')(x_in)
    return heatmaps


def Classifier():
    """Classification between single and interacting hands branch
    """
    def classifier(x_in):
        x = inputs = tf.keras.layers.Input(x_in.shape[1:])
        x = tf.keras.layers.Conv2D(256, 1, strides=1)(x)
        x = tf.keras.layers.LeakyReLU(0.1)(x)
        x = tf.keras.layers.MaxPool2D((2, 2))(x)
        x = tf.keras.layers.Flatten()(x)
        x = tf.keras.layers.Dense(128, kernel_regularizer='l2')(x)
        x = tf.keras.layers.LeakyReLU(0.1)(x)
        x = outputs = tf.keras.layers.Dense(1, activation='sigmoid')(x)
        return tf.keras.Model(inputs=inputs, outputs=outputs, name='InteractingClassifier')(x_in)
    return classifier


class CustomModel(tf.keras.Model):
    """Custom tf.keras.Model
    """
    def __init__(self, loss_weights, **kwargs):
        """
        Args:
            loss_weights (tuple): classification weights
        """
        super().__init__(**kwargs)
        # overall loss averaging metric
        self.avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32)
        # loss averaging metric for each model branch
        self.handness_loss = tf.keras.metrics.Mean('handness_loss', dtype=tf.float32)
        self.single_hand_loss = tf.keras.metrics.Mean('single_hand_loss', dtype=tf.float32)
        self.interacting_loss = tf.keras.metrics.Mean('interacting_loss', dtype=tf.float32)
        # loss function callable
        self.loss = HandPoseLoss(*loss_weights)
    
    def train_step(self, data):
        """Training step

        Args:
            data (tuple): batched input images and a dictionary of labels for each branch

        Returns:
            dict: loss results
        """
        images, labels = data
        # wrap loss in autodiff
        with tf.GradientTape() as tape:
            outputs = self(images, training=True)
            reg_loss = tf.reduce_sum(self.losses)
            handness_loss, single_hand_loss, interacting_loss = self.loss(labels, outputs)
            pred_loss = handness_loss + single_hand_loss + interacting_loss
            total_loss = tf.reduce_sum(pred_loss) + reg_loss
        # calculate gradients and apply them with optimizer
        gradients = tape.gradient(total_loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
        # update loss metrics and output them
        self.avg_loss.update_state(total_loss)
        self.handness_loss.update_state(tf.reduce_sum(handness_loss))
        self.single_hand_loss.update_state(tf.reduce_sum(single_hand_loss))
        self.interacting_loss.update_state(tf.reduce_sum(interacting_loss))
        return {'loss': self.avg_loss.result(), 'handness_loss': self.handness_loss.result(),
                'single_hand_loss': self.single_hand_loss.result(), 'interacting_loss': self.interacting_loss.result()}
    
    def test_step(self, data):
        """Testing step

        Args:
            data (tuple): batched input images and a dictionary of labels for each branch

        Returns:
            dict: loss results
        """
        images, labels = data
        outputs = self(images, training=True)
        handness_loss, single_hand_loss, interacting_loss = self.loss(labels, outputs)
        pred_loss = handness_loss + single_hand_loss + interacting_loss
        total_loss = tf.reduce_sum(pred_loss)
        self.avg_loss.update_state(total_loss)
        self.handness_loss.update_state(tf.reduce_sum(handness_loss))
        self.single_hand_loss.update_state(tf.reduce_sum(single_hand_loss))
        self.interacting_loss.update_state(tf.reduce_sum(interacting_loss))
        return {'loss': self.avg_loss.result(), 'handness_loss': self.handness_loss.result(),
                'single_hand_loss': self.single_hand_loss.result(), 'interacting_loss': self.interacting_loss.result()}
    
    @property
    def metrics(self):
        """Define tracked metrics throughout iterations as property. This way the metrics are properly managed by the model internally.

        Returns:
            list: list of metrics
        """
        return [self.avg_loss, self.handness_loss, self.single_hand_loss, self.interacting_loss]


def HandKeypointEstimator(backbone, input_size, hm_dims=None, imagenet=False, training=False, decode='argmax', delta=4, loss_weights=(1., 1.)):
    """Hand pose estimation model

    Args:
        backbone (str): efficientnet type backbone, one of {efficientnet-b0, efficientnet-b1, efficientnet-b2, efficientnet-b3}
        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.
        loss_weights (tuple, optional): classification weights for [single, interacting] classes. Defaults to (1., 1.).

    Returns:
        tf.keras.Model: constructed hand pose estimation 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,
        'efficientnet-b3': tf.keras.applications.efficientnet.EfficientNetB3,
    }
    x = inputs = tf.keras.layers.Input((*input_size, 3), name='input')
    x = backbone_dict[backbone](include_top=False, weights='imagenet' if imagenet else None)(x)
    x_features = tf.keras.layers.Dropout(0.3, name='dropout')(x)
    x_single_hand = SingleHandHeatmaps()(x_features)
    x_interacting = InteractingHeatmaps()(x_features)
    x_handness = Classifier()(x_features)
    # if training is enabled or decoding is disabled, return raw heatmaps
    if training or decode is None:
        outputs = [x_handness, x_single_hand, x_interacting]
    else:
        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_keypoints_single_hand = tf.keras.layers.Lambda(lambda x: decode_method[decode](x) / hm_dims)(x_single_hand)
        x_keypoints_interacting = tf.keras.layers.Lambda(lambda x: decode_method[decode](x) / hm_dims)(x_interacting)
        outputs = [x_handness, x_keypoints_single_hand, x_keypoints_interacting]
    return CustomModel(loss_weights, inputs=inputs, outputs=outputs, name='HandKeypointEstimator')


def HandPoseLoss(negative_weight, positive_weight):
    """Hand pose estimation loss function

    Args:
        negative_weight (float): single hand loss weight
        positive_weight (float): interacting hands loss weight
        
    Returns:
        callable: callable loss function loss(y_true, y_pred)
    """
    def hand_pose_loss(y_true, y_pred):
        # classification loss - weighted cross entropy
        handness_true = tf.cast(y_true['class'], tf.float32)
        weight = handness_true * positive_weight + (1. - handness_true) * negative_weight
        handness_loss = tf.keras.losses.binary_crossentropy(y_true['class'], y_pred[0]) * weight
        handness_loss = tf.reduce_mean(handness_loss)
        # heatmap losses - mean squared error, loss is masked by keypoint validity
        single_hand_mask = tf.abs(tf.cast(tf.reshape(y_true['class'], (-1, 1, 1, 1)), tf.float32) - 1)
        interacting_mask = tf.cast(tf.reshape(y_true['class'], (-1, 1, 1, 1)), tf.float32)
        validity_mask = tf.cast(tf.reshape(y_true['validity'], (-1, 1, 1, 42)), tf.float32)

        single_hand_loss = (tf.square(y_true['heatmaps'][..., :21] - y_pred[1]) * single_hand_mask) * validity_mask[..., :21] * negative_weight
        single_hand_loss = tf.reduce_mean(single_hand_loss)

        interacting_loss = (tf.square(y_true['heatmaps'] - y_pred[2]) * interacting_mask) * validity_mask * positive_weight
        interacting_loss = tf.reduce_mean(interacting_loss)
        return [handness_loss, single_hand_loss, interacting_loss]
    return hand_pose_loss


def load_model(weights_path, *args, **kwargs):
    """Load hand 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 hand pose estimation model
    """
    # construct model in training mode and load its weights
    model = HandKeypointEstimator(*args, **kwargs, training=True)
    model.load_weights(weights_path).expect_partial()
    # construct inference mode model
    inf_model = HandKeypointEstimator(*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
