"""
DeepLabCut2.2 Toolbox (deeplabcut.org)
© A. & M. Mathis Labs
https://github.com/AlexEMG/DeepLabCut

Please see AUTHORS for contributors.
https://github.com/AlexEMG/DeepLabCut/blob/master/AUTHORS
Licensed under GNU Lesser General Public License v3.0

Adapted from DeeperCut by Eldar Insafutdinov
https://github.com/eldar/pose-tensorflow
"""

import re
import functools

import tensorflow as tf
import tensorflow.contrib.slim as slim
from tensorflow.contrib.slim.nets import resnet_v1

from deeplabcut.pose_estimation_tensorflow.dataset.pose_dataset import Batch
import deeplabcut.pose_estimation_tensorflow.nnet.efficientnet_builder as eff
from deeplabcut.pose_estimation_tensorflow.nnet import mobilenet_v2, mobilenet, conv_blocks
from deeplabcut.pose_estimation_tensorflow.nnet import losses

vers = (tf.__version__).split(".")
if int(vers[0]) == 1 and int(vers[1]) > 12:
    TF = tf.compat.v1
else:
    TF = tf

# Change the stride from 2 to 1 to get 16x downscaling instead of 32x.
mobilenet_v2.V2_DEF["spec"][14] = mobilenet.op(conv_blocks.expanded_conv, stride=1, num_outputs=160)
def wrapped_partial(func, *args, **kwargs):
    partial_func = functools.partial(func, *args, **kwargs)
    functools.update_wrapper(partial_func, func)
    return partial_func

net_funcs = {
    "resnet_50": resnet_v1.resnet_v1_50,
    "resnet_101": resnet_v1.resnet_v1_101,
    "resnet_152": resnet_v1.resnet_v1_152,
    'mobilenet_v2_1.0': mobilenet_v2.mobilenet_base,
    'mobilenet_v2_0.75': wrapped_partial(mobilenet_v2.mobilenet_base,
                                        depth_multiplier=0.75,
                                        final_endpoint="layer_19",
                                        finegrain_classification_mode=True),
    'mobilenet_v2_0.5': wrapped_partial(mobilenet_v2.mobilenet_base,
                                        depth_multiplier=0.5,
                                        final_endpoint="layer_19",
                                        finegrain_classification_mode=True),
    'mobilenet_v2_0.35': wrapped_partial(mobilenet_v2.mobilenet_base,
                                        depth_multiplier=0.35,
                                        final_endpoint="layer_19",
                                        finegrain_classification_mode=True),
    'mobilenet_v2_0.1': wrapped_partial(mobilenet_v2.mobilenet_base,
                                        depth_multiplier=0.1,
                                        final_endpoint="layer_19",
                                        finegrain_classification_mode=True),
    'mobilenet_v2_0.35_10': wrapped_partial(mobilenet_v2.mobilenet_base,
                                            depth_multiplier=0.35,
                                            final_endpoint="layer_10",
                                            finegrain_classification_mode=True),
    'mobilenet_v2_0.1_10':  wrapped_partial(mobilenet_v2.mobilenet_base,
                                            depth_multiplier=0.1,
                                            final_endpoint="layer_10",
                                            finegrain_classification_mode=True)
}

#https://towardsdatascience.com/complete-architectural-details-of-all-efficientnet-models-5fd5b736142
parallel_layers = {
    "b0": "4",
    "b1": "7",
    "b2": "7",
    "b3": "7",
    "b4": "9",
    "b5": "12",
    "b6": "14",
    "b7": "17"
    }


def prediction_layer(cfg, input, name, num_outputs):
    with slim.arg_scope(
        [slim.conv2d, slim.conv2d_transpose],
        padding="SAME",
        activation_fn=None,
        normalizer_fn=None,
        weights_regularizer=slim.l2_regularizer(cfg['weight_decay']),
    ):
        with tf.variable_scope(name):
            pred = slim.conv2d_transpose(
                input, num_outputs, kernel_size=[3, 3], stride=2, scope="block4"
            )
            return pred


def get_batch_spec(cfg):
    num_joints = cfg['num_joints']
    batch_size = cfg['batch_size']
    batch_spec = {
        Batch.inputs: [batch_size, None, None, 3],
        Batch.part_score_targets: [
            batch_size,
            None,
            None,
            num_joints + cfg.get("num_idchannel", 0),
        ],
        Batch.part_score_weights: [
            batch_size,
            None,
            None,
            num_joints + cfg.get("num_idchannel", 0),
        ],
    }
    if cfg['location_refinement']:
        batch_spec[Batch.locref_targets] = [batch_size, None, None, num_joints * 2]
        batch_spec[Batch.locref_mask] = [batch_size, None, None, num_joints * 2]
    if cfg['pairwise_predict']:
        print("Getting specs", cfg['dataset_type'], cfg['num_limbs'], cfg['num_joints'])
        if (
            "multi-animal" not in cfg['dataset_type']
        ):  # this can be used for pairwise conditional
            batch_spec[Batch.pairwise_targets] = [
                batch_size,
                None,
                None,
                num_joints * (num_joints - 1) * 2,
            ]
            batch_spec[Batch.pairwise_mask] = [
                batch_size,
                None,
                None,
                num_joints * (num_joints - 1) * 2,
            ]
        else:  # train partaffinity fields
            batch_spec[Batch.pairwise_targets] = [
                batch_size,
                None,
                None,
                cfg['num_limbs'] * 2,
            ]
            batch_spec[Batch.pairwise_mask] = [
                batch_size,
                None,
                None,
                cfg['num_limbs'] * 2,
            ]
    return batch_spec


class PoseNet:
    def __init__(self, cfg):
        self.cfg = cfg
        if 'use_batch_norm' not in self.cfg.keys():
            self.cfg['use_batch_norm'] = False
        if 'use_drop_out' not in self.cfg.keys():
            self.cfg['use_drop_out'] = False

    def extract_features(self, inputs):
        mean = tf.constant(
            self.cfg['mean_pixel'], dtype=tf.float32, shape=[1, 1, 1, 3], name="img_mean"
        )
        im_centered = inputs - mean

        if 'resnet' in self.cfg['net_type']:
            # The next part of the code depends upon which tensorflow version you have.
            vers = tf.__version__
            vers = vers.split(
                "."
            )  # Updated based on https://github.com/AlexEMG/DeepLabCut/issues/44

            net_fun = net_funcs[self.cfg['net_type']]
            if int(vers[0]) == 1 and int(vers[1]) < 4:  # check if lower than version 1.4.
                with slim.arg_scope(resnet_v1.resnet_arg_scope(False)):
                    net, end_points = net_fun(
                        im_centered, global_pool=False, output_stride=16
                    )
            else:
                with slim.arg_scope(resnet_v1.resnet_arg_scope()):
                    net, end_points = net_fun(
                        im_centered, global_pool=False, output_stride=16, is_training=False
                    )
        elif 'mobilenet' in self.cfg['net_type']:
            net_fun = net_funcs[self.cfg['net_type']]
            with slim.arg_scope(mobilenet_v2.training_scope()):
                net, end_points = net_fun(im_centered)
        elif 'efficientnet' in self.cfg['net_type']:
            im_centered /= tf.constant(eff.STDDEV_RGB, shape=[1, 1, 3])
            net, end_points = eff.build_model_base(im_centered,
                                                   self.cfg['net_type'],
                                                   use_batch_norm=self.cfg['use_batch_norm'],
                                                   drop_out=self.cfg['use_drop_out'])
        return net, end_points

    def prediction_layers(
        self,
        features,
        end_points,
        input_shape,
        reuse=None,
        is_training=False,
        scope="pose",
    ):
        cfg = self.cfg
        if "resnet" in cfg['net_type']:
            num_layers = re.findall("resnet_([0-9]*)", cfg['net_type'])[0]
            layer_name = (
                "resnet_v1_{}".format(num_layers) + "/block{}/unit_{}/bottleneck_v1"
            )
            mid_pt = layer_name.format(2,3)
        elif "mobilenet" in cfg['net_type']:
            mid_pt = "layer_7"
        elif "efficientnet" in cfg['net_type']:
            mid_pt = "block_"+parallel_layers[cfg['net_type'].split('-')[1]]

        final_dims = tf.ceil(
            tf.divide(input_shape[1:3], tf.convert_to_tensor(16))
        )
        interim_dims = tf.scalar_mul(2, final_dims)
        interim_dims = tf.cast(interim_dims, tf.int32)
        bank_3 = end_points[mid_pt]
        bank_3 = tf.image.resize_images(bank_3, interim_dims)

        with slim.arg_scope(
            [slim.conv2d],
            padding="SAME",
            normalizer_fn=None,
            weights_regularizer=slim.l2_regularizer(cfg['weight_decay']),
        ):
            with tf.variable_scope("decoder_filters"):
                bank_3 = slim.conv2d(bank_3, cfg['bank3'], 1, scope="decoder_parallel_1")

        with slim.arg_scope(
            [slim.conv2d_transpose],
            padding="SAME",
            normalizer_fn=None,
            weights_regularizer=slim.l2_regularizer(cfg['weight_decay']),
        ):
            with tf.variable_scope("upsampled_features"):
                upsampled_features = slim.conv2d_transpose(
                    features, cfg['bank5'], kernel_size=[3, 3], stride=2, scope="block4"
                )
        net = tf.concat([bank_3, upsampled_features], 3)

        out = {}
        with tf.variable_scope(scope, reuse=reuse):
            out["part_pred"] = prediction_layer(
                cfg, net, "part_pred", cfg['num_joints'] + cfg.get("num_idchannel", 0)
            )
            if cfg['location_refinement']:
                out["locref"] = prediction_layer(
                    cfg, net, "locref_pred", cfg['num_joints'] * 2
                )
            if cfg['pairwise_predict'] and "multi-animal" not in cfg['dataset_type']:
                out["pairwise_pred"] = prediction_layer(
                    cfg, net, "pairwise_pred", cfg['num_joints'] * (cfg['num_joints'] - 1) * 2
                )
            if cfg['partaffinityfield_predict'] and "multi-animal" in cfg['dataset_type']:
                out["pairwise_pred"] = prediction_layer(
                    cfg, net, "pairwise_pred", cfg['num_limbs'] * 2
                )

            if cfg['intermediate_supervision'] and "efficientnet" not in cfg['net_type']:
                if "mobilenet" in cfg['net_type']:
                    out["part_pred_interm"] = prediction_layer(
                        cfg,
                        end_points["layer_" + str(cfg["intermediate_supervision_layer"])],
                        "intermediate_supervision",
                        cfg['num_joints'],
                    )
                elif "resnet" in cfg['net_type']:
                    interm_name = layer_name.format(3, cfg['intermediate_supervision_layer'])
                    block_interm_out = end_points[interm_name]
                    out["part_pred_interm"] = prediction_layer(
                        cfg,
                        block_interm_out,
                        "intermediate_supervision",
                        cfg['num_joints'] + cfg.get("num_idchannel", 0),
                    )

        return out

    def get_net(self, inputs):
        net, end_points = self.extract_features(inputs)
        # return self.prediction_layers(net, end_points)
        return self.prediction_layers(net, end_points, tf.shape(inputs))

    def test(self, inputs):
        heads = self.get_net(inputs)
        return self.add_inference_layers(heads)

    def add_inference_layers(self, heads):
        """ initialized during inference """
        prob = tf.sigmoid(heads["part_pred"])
        outputs = {"part_prob": prob}
        if self.cfg['location_refinement']:
            outputs["locref"] = heads["locref"]
        if self.cfg['pairwise_predict'] or self.cfg['partaffinityfield_predict']:
            outputs["pairwise_pred"] = heads["pairwise_pred"]
        return outputs

    def train(self, batch):
        cfg = self.cfg

        heads = self.get_net(batch[Batch.inputs])
        weigh_part_predictions = cfg['weigh_part_predictions']
        part_score_weights = (
            batch[Batch.part_score_weights] if weigh_part_predictions else 1.0
        )

        def add_part_loss(pred_layer):
            return TF.losses.sigmoid_cross_entropy(
                batch[Batch.part_score_targets], heads[pred_layer], part_score_weights
            )

        loss = {}
        loss["part_loss"] = add_part_loss("part_pred")
        total_loss = loss["part_loss"]
        if cfg['intermediate_supervision'] and "efficientnet" not in cfg['net_type']:
            loss["part_loss_interm"] = add_part_loss("part_pred_interm")
            total_loss = total_loss + loss["part_loss_interm"]

        if cfg['location_refinement']:
            locref_pred = heads["locref"]
            locref_targets = batch[Batch.locref_targets]
            locref_weights = batch[Batch.locref_mask]
            loss_func = (
                losses.huber_loss
                if cfg['locref_huber_loss']
                else tf.losses.mean_squared_error
            )
            loss["locref_loss"] = cfg['locref_loss_weight'] * loss_func(
                locref_targets, locref_pred, locref_weights
            )
            total_loss = total_loss + loss["locref_loss"]

        if cfg['pairwise_predict'] or cfg['partaffinityfield_predict']:
            "setting pw bodypart loss..."
            pairwise_pred = heads["pairwise_pred"]
            pairwise_targets = batch[Batch.pairwise_targets]
            pairwise_weights = batch[Batch.pairwise_mask]

            loss_func = (
                losses.huber_loss
                if cfg['pairwise_huber_loss']
                else tf.losses.mean_squared_error
            )
            loss["pairwise_loss"] = cfg['pairwise_loss_weight'] * loss_func(
                pairwise_targets, pairwise_pred, pairwise_weights
            )
            total_loss = total_loss + loss["pairwise_loss"]

        # loss['total_loss'] = slim.losses.get_total_loss(add_regularization_losses=params.regularize)
        loss["total_loss"] = total_loss
        return loss
