"""
Interhand model implementation
"""
import tensorflow as tf

from .estimation_utils import soft_argmax, decode_keypoints_argmax_3d


def ResidualBlock(x, filters, stride, name, expand=False):
    """Backbone residual block
    """
    inputs = x
    if expand:
        inputs = tf.keras.layers.Conv2D(4*filters, 1, strides=stride, use_bias=False, name=name+'downsample.0')(inputs)
        inputs = tf.keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.1, name=name+'downsample.1')(inputs)
    x = tf.keras.layers.Conv2D(filters, 1, strides=1, use_bias=False, name=name+'conv1')(x)
    x = tf.keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.1, name=name+'bn1')(x)
    x = tf.keras.layers.ReLU()(x)
    x = tf.keras.layers.ZeroPadding2D((1, 1))(x)
    x = tf.keras.layers.Conv2D(filters, 3, strides=stride, use_bias=False, name=name+'conv2')(x)
    x = tf.keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.1, name=name+'bn2')(x)
    x = tf.keras.layers.ReLU()(x)
    x = tf.keras.layers.Conv2D(4*filters, 1, strides=1, use_bias=False, name=name+'conv3')(x)
    x = tf.keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.1, name=name+'bn3')(x)
    x = tf.keras.layers.Add()([inputs, x])
    x = tf.keras.layers.ReLU()(x)
    return x


def ResNetBackbone(name=None):
    """Resnet backbone model
    """
    filters = 64
    x = inputs = tf.keras.layers.Input((None, None, 3))
    x = tf.keras.layers.ZeroPadding2D((3, 3))(x)
    x = tf.keras.layers.Conv2D(filters, 7, strides=2, use_bias=False, name='resnet.conv1')(x)
    x = tf.keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.1, name='resnet.bn1')(x)
    x = tf.keras.layers.ReLU()(x)
    x = tf.keras.layers.ZeroPadding2D(padding=(1, 1))(x)
    x = tf.keras.layers.MaxPool2D(pool_size=(3, 3), strides=2)(x)
    blocks = [3, 4, 6, 3]
    stride = 1
    for i, i_blocks in enumerate(blocks):
        x = ResidualBlock(x, filters, stride, f'resnet.layer{i+1}.0.', True)
        for j in range(i_blocks-1):
            x = ResidualBlock(x, filters, 1, f'resnet.layer{i+1}.{j+1}.')
        filters *= 2
        stride = 2
    outputs = x
    return tf.keras.Model(inputs=inputs, outputs=outputs, name=name)


def SinglePoseNet(joints_hm_dims, idx):
    """Single hand pose estimation branch
    """
    def single_pose_net(x):
        name = f'pose_net.joint_deconv_{idx}.'
        x = tf.keras.layers.Conv2DTranspose(256, 4, strides=2, padding='same', use_bias=False, name=name+'0')(x)
        x = tf.keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.1, name=name+'1')(x)
        x = tf.keras.layers.ReLU()(x)
        x = tf.keras.layers.Conv2DTranspose(256, 4, strides=2, padding='same', use_bias=False, name=name+'3')(x)
        x = tf.keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.1, name=name+'4')(x)
        x = tf.keras.layers.ReLU()(x)
        x = tf.keras.layers.Conv2DTranspose(256, 4, strides=2, padding='same', use_bias=False, name=name+'6')(x)
        x = tf.keras.layers.BatchNormalization(epsilon=1e-5, momentum=0.1, name=name+'7')(x)
        x = tf.keras.layers.ReLU()(x)
        x = tf.keras.layers.Conv2D(joints_hm_dims[2]*21, 1, strides=1, name=f'pose_net.joint_conv_{idx}.0')(x)
        x = tf.keras.layers.Reshape((joints_hm_dims[0], joints_hm_dims[1], 21, joints_hm_dims[2]))(x)
        return x
    return single_pose_net


def PoseNet(joints_hm_dims, name=None):
    """Right and left hand pose estimation branch
    """
    def pose_net(x_in):
        x = inputs = tf.keras.layers.Input(x_in.shape[1:])
        x1 = SinglePoseNet(joints_hm_dims, 1)(x)
        x2 = SinglePoseNet(joints_hm_dims, 2)(x)
        x = tf.keras.layers.Concatenate(axis=3, name='heatmaps')([x1, x2])
        outputs = x = tf.keras.layers.Permute((1, 2, 4, 3))(x)
        return tf.keras.Model(inputs=inputs, outputs=outputs, name=name)(x_in)
    return pose_net


def HandNet(name=None):
    """Handedness classification branch
    """
    def hand_net(x_in):
        x = inputs = tf.keras.layers.Input(x_in.shape[1:])
        x = tf.keras.layers.Dense(512, name='pose_net.hand_fc.0')(x)
        x = tf.keras.layers.ReLU()(x)
        outputs = x = tf.keras.layers.Dense(2, activation='sigmoid', name='pose_net.hand_fc.2')(x)
        return tf.keras.Model(inputs=inputs, outputs=outputs, name=name)(x_in)
    return hand_net


def RelativeDepthNet(root_hm_dim, name=None):
    """Relative root depth branch
    """
    def relative_depth_net(x_in):
        x = inputs = tf.keras.layers.Input(x_in.shape[1:])
        x = tf.keras.layers.Dense(512, name='pose_net.root_fc.0')(x)
        x = tf.keras.layers.ReLU()(x)
        x = tf.keras.layers.Dense(root_hm_dim, name='pose_net.root_fc.2')(x)
        outputs = x = tf.keras.layers.Lambda(lambda z: soft_argmax(z), name='relative_root_depth')(x)
        return tf.keras.Model(inputs=inputs, outputs=outputs, name=name)(x_in)
    return relative_depth_net


def transform_pose(y_pose, joints_hm_dims, depth_dim):
    """Scale predicted keypoints to <0, 1> range

    Args:
        y_pose (tf.Tensor): predicted keypoints
        joints_hm_dims (tuple): size of generated heatmaps
        depth_dim (int): depth representation vector length

    Returns:
        tf.Tensor: transformed keypoints
    """
    # scale 2D keypoints
    y_pose_2d = y_pose[:, :, :2] / joints_hm_dims[:2]
    # scale depth dimension
    y_pose_depth = (y_pose[:, :, 2:] / joints_hm_dims[2] * 2 - 1) * (depth_dim / 2)
    y_pose = tf.concat((y_pose_2d, y_pose_depth), axis=2)
    return y_pose


def transform_relative_depth(y_depth, root_hm_dim, depth_dim):
    """Scale predicted relative root depth

    Args:
        y_depth (tf.Tensor): predicted depth
        root_hm_dim (int): relative hand depth dimension
        depth_dim (int): depth representation vector length

    Returns:
        tf.Tensor: transformed depth
    """
    return (y_depth / root_hm_dim * 2 - 1) * (depth_dim / 2)


class CustomModel(tf.keras.Model):
    """Custom tf.keras.Model
    """
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32)
        self.loss = HandPoseLoss()
    
    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)
            pred_loss = self.loss(labels, outputs)
            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 metric and output it
        self.avg_loss.update_state(total_loss)
        return {"loss": self.avg_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]


def InterhandEstimator(input_size, joints_hm_dims, depth_dim, root_hm_dim, training=False):
    """Interhand 3D hand pose estimation model

    Args:
        input_size (tuple): input image size in first two dimensions
        joints_hm_dims (tuple): size of generated heatmaps
        depth_dim (int): depth representation vector length
        root_hm_dim (int): relative hand depth dimension
        training (bool, optional): training flag, if False - testing mode is enabled. Defaults to False.

    Returns:
        tf.keras.Model: constructed hand pose estimation model
    """
    x = inputs = tf.keras.layers.Input((input_size[0], input_size[1], 3))
    x = ResNetBackbone(name='ResnetBackbone')(x)
    y_pose = PoseNet(joints_hm_dims, name='PoseNet')(x)
    x = tf.keras.layers.GlobalAveragePooling2D()(x)
    y_handness = HandNet(name='HandNet')(x)
    y_depth = RelativeDepthNet(root_hm_dim, name='DepthNet')(x)
    # postprocess heatmaps if testing
    if not training:
        y_pose = tf.keras.layers.Lambda(lambda x: decode_keypoints_argmax_3d(x))(y_pose)
        y_pose = tf.keras.layers.Lambda(lambda x: transform_pose(x, joints_hm_dims, depth_dim))(y_pose)
        y_depth = tf.keras.layers.Lambda(lambda x: transform_relative_depth(x, root_hm_dim, depth_dim))(y_depth)
    outputs = [y_handness, y_pose, y_depth]
    return CustomModel(inputs=inputs, outputs=outputs, name='InterhandEstimator')


def HandPoseLoss():
    """Interhand pose estimation loss function
    
    Returns:
        callable: loss function loss(y_true, y_pred)
    """
    def hand_pose_loss(y_true, y_pred):
        # classification loss
        handness_loss = tf.keras.losses.binary_crossentropy(y_true['handness'], y_pred[0]) * y_true['handness_valid']
        # heatmap loss
        pose_loss = tf.square(y_true['heatmaps'] - y_pred[1]) * tf.reshape(y_true['joints_valid'], (-1, 1, 1, 42, 1))
        pose_loss = tf.reduce_sum(pose_loss, axis=(1, 2, 3, 4))
        # depth regression loss
        depth_loss = tf.abs(y_true['relative_root_depth'] - y_pred[2]) * y_true['relative_root_valid']
        depth_loss = tf.reduce_sum(depth_loss, axis=1)
        return handness_loss + pose_loss + depth_loss
    return hand_pose_loss
