import tensorflow as tf

import numpy as np

import part_detector

import settings

import utils

import os

from abc import abstractmethod

from functools import lru_cache

from scipy.stats import norm

from inception_resnet_v2 import inception_resnet_v2_arg_scope, inception_resnet_v2

import tensorflow.contrib.layers as layers

slim = tf.contrib.slim

SUMMARY_PATH = settings.LOGDIR_PATH

KEY_SUMMARIES = tf.GraphKeys.SUMMARIES

KEY_SUMMARIES_PER_JOINT = ['summary_joint_%02d' % i for i in range(16)]

class HumanPoseNN(object):

    """

    The neural network used for pose estimation.

    """

    def __init__(self, log_name, heatmap_size, image_size, loss_type = 'SCE', is_training = True):

        tf.set_random_seed(0)

        if loss_type not in { 'MSE', 'SCE' }:

            raise NotImplementedError('Loss function should be either MSE or SCE!')

        self.log_name = log_name

        self.heatmap_size = heatmap_size

        self.image_size = image_size

        self.is_train = is_training

        self.loss_type = loss_type

        # Initialize placeholders

        self.input_tensor = tf.placeholder(

            dtype = tf.float32,

            shape = (None, image_size, image_size, 3),

            name = 'input_image')

        self.present_joints = tf.placeholder(

            dtype = tf.float32,

            shape = (None, 16),

            name = 'present_joints')

        self.inside_box_joints = tf.placeholder(

            dtype = tf.float32,

            shape = (None, 16),

            name = 'inside_box_joints')

        self.desired_heatmap = tf.placeholder(

            dtype = tf.float32,

            shape = (None, heatmap_size, heatmap_size, 16),

            name = 'desired_heatmap')

        self.desired_points = tf.placeholder(

            dtype = tf.float32,

            shape = (None, 2, 16),

            name = 'desired_points')

        self.network = self.pre_process(self.input_tensor)

        self.network, self.feature_tensor = self.get_network(self.network, is_training)

        self.sigm_network = tf.sigmoid(self.network)

        self.smoothed_sigm_network = self._get_gauss_smoothing_net(self.sigm_network, std = 0.7)

        self.loss_err = self._get_loss_function(loss_type)

        self.euclidean_dist = self._euclidean_dist_err()

        self.euclidean_dist_per_joint = self._euclidean_dist_per_joint_err()

        if is_training:

            self.global_step = tf.Variable(0, name = 'global_step', trainable = False)

            self.learning_rate = tf.placeholder(

                dtype = tf.float32,

                shape = [],

                name = 'learning_rate')

            self.optimize = layers.optimize_loss(loss = self.loss_err,

                                                 global_step = self.global_step,

                                                 learning_rate = self.learning_rate,

                                                 optimizer = tf.train.RMSPropOptimizer(self.learning_rate),

                                                 clip_gradients = 2.0

                                                 )

        self.sess = tf.Session()

        self.sess.run(tf.global_variables_initializer())

        if log_name is not None:

            self._init_summaries()

    def _init_summaries(self):

        if self.is_train:

            logdir = os.path.join(SUMMARY_PATH, self.log_name, 'train')

            self.summary_writer = tf.summary.FileWriter(logdir)

            self.summary_writer_by_points = [tf.summary.FileWriter(os.path.join(logdir, 'point_%02d' % i))

                                             for i in range(16)]

            tf.scalar_summary('Average euclidean distance', self.euclidean_dist, collections = [KEY_SUMMARIES])

            for i in range(16):

                tf.scalar_summary('Joint euclidean distance', self.euclidean_dist_per_joint[i],

                                  collections = [KEY_SUMMARIES_PER_JOINT[i]])

            self.create_summary_from_weights()

            self.ALL_SUMMARIES = tf.merge_all_summaries(KEY_SUMMARIES)

            self.SUMMARIES_PER_JOINT = [tf.merge_all_summaries(KEY_SUMMARIES_PER_JOINT[i]) for i in range(16)]

        else:

            logdir = os.path.join(SUMMARY_PATH, self.log_name, 'test')

            self.summary_writer = tf.summary.FileWriter(logdir)

    def _get_loss_function(self, loss_type):

        loss_dict = {

            'MSE': self._loss_mse(),

            'SCE': self._loss_cross_entropy()

        }

        return loss_dict[loss_type]

    @staticmethod

    @lru_cache()

    def _get_gauss_filter(size = 15, std = 1.0, kernel_sum = 1.0):

        samples = norm.pdf(np.linspace(-2, 2, size), 0, std)

        samples /= np.sum(samples)

        samples *= kernel_sum ** 0.5

        samples = np.expand_dims(samples, 0)

        weights = np.zeros(shape = (1, size, 16, 1), dtype = np.float32)

        for i in range(16):

            weights[:, :, i, 0] = samples

        return weights

    @staticmethod

    def _get_gauss_smoothing_net(net, size = 15, std = 1.0, kernel_sum = 1.0):

        filter_h = HumanPoseNN._get_gauss_filter(size, std, kernel_sum)

        filter_v = filter_h.swapaxes(0, 1)

        net = tf.nn.depthwise_conv2d(net, filter = filter_h, strides = [1, 1, 1, 1], padding = 'SAME',

                                     name = 'SmoothingHorizontal')

        net = tf.nn.depthwise_conv2d(net, filter = filter_v, strides = [1, 1, 1, 1], padding = 'SAME',

                                     name = 'SmoothingVertical')

        return net

    def generate_output(self, shape, presented_parts, labels, sigma):

        heatmap_dict = {

            'MSE': utils.get_gauss_heat_map(

                shape = shape, is_present = presented_parts,

                mean = labels, sigma = sigma),

            'SCE': utils.get_binary_heat_map(

                shape = shape, is_present = presented_parts,

                centers = labels, diameter = sigma)

        }

        return heatmap_dict[self.loss_type]

    def _adjust_loss(self, loss_err):

        # Shape: [batch, joints]

        loss = tf.reduce_sum(loss_err, [1, 2])

        # Stop error propagation of joints that are not presented

        loss = tf.multiply(loss, self.present_joints)

        # Compute average loss of presented joints

        num_of_visible_joints = tf.reduce_sum(self.present_joints)

        loss = tf.reduce_sum(loss) / num_of_visible_joints

        return loss

    def _loss_mse(self):

        sq = tf.squared_difference(self.sigm_network, self.desired_heatmap)

        loss = self._adjust_loss(sq)

        return loss

    def _loss_cross_entropy(self):

        ce = tf.nn.sigmoid_cross_entropy_with_logits(logits = self.network, labels = self.desired_heatmap)

        loss = self._adjust_loss(ce)

        return loss

    def _joint_highest_activations(self):

        highest_activation = tf.reduce_max(self.smoothed_sigm_network, [1, 2])

        return highest_activation

    def _joint_positions(self):

        highest_activation = tf.reduce_max(self.sigm_network, [1, 2])

        x = tf.argmax(tf.reduce_max(self.smoothed_sigm_network, 1), 1)

        y = tf.argmax(tf.reduce_max(self.smoothed_sigm_network, 2), 1)

        x = tf.cast(x, tf.float32)

        y = tf.cast(y, tf.float32)

        a = tf.cast(highest_activation, tf.float32)

        scale_coef = (self.image_size / self.heatmap_size)

        x *= scale_coef

        y *= scale_coef

        out = tf.stack([y, x, a])

        return out

    def _euclidean_dist_err(self):

        # Work only with joints that are presented inside frame

        l2_dist = tf.multiply(self.euclidean_distance(), self.inside_box_joints)

        # Compute average loss of presented joints

        num_of_visible_joints = tf.reduce_sum(self.inside_box_joints)

        l2_dist = tf.reduce_sum(l2_dist) / num_of_visible_joints

        return l2_dist

    def _euclidean_dist_per_joint_err(self):

        # Work only with joints that are presented inside frame

        l2_dist = tf.multiply(self.euclidean_distance(), self.inside_box_joints)

        # Average euclidean distance of presented joints

        present_joints = tf.reduce_sum(self.inside_box_joints, 0)

        err = tf.reduce_sum(l2_dist, 0) / present_joints

        return err

    def _restore(self, checkpoint_path, variables):

        saver = tf.train.Saver(variables)

        saver.restore(self.sess, checkpoint_path)

    def _save(self, checkpoint_path, name, variables):

        if not os.path.exists(checkpoint_path):

            os.mkdir(checkpoint_path)

        checkpoint_name_path = os.path.join(checkpoint_path, '%s.ckpt' % name)

        saver = tf.train.Saver(variables)

        saver.save(self.sess, checkpoint_name_path)

    def euclidean_distance(self):

        x = tf.argmax(tf.reduce_max(self.smoothed_sigm_network, 1), 1)

        y = tf.argmax(tf.reduce_max(self.smoothed_sigm_network, 2), 1)

        x = tf.cast(x, tf.float32)

        y = tf.cast(y, tf.float32)

        dy = tf.squeeze(self.desired_points[:, 0, :])

        dx = tf.squeeze(self.desired_points[:, 1, :])

        sx = tf.squared_difference(x, dx)

        sy = tf.squared_difference(y, dy)

        l2_dist = tf.sqrt(sx + sy)

        return l2_dist

    def feed_forward(self, x):

        out = self.sess.run(self.sigm_network, feed_dict = {

            self.input_tensor: x

        })

        return out

    def heat_maps(self, x):

        out = self.sess.run(self.smoothed_sigm_network, feed_dict = {

            self.input_tensor: x

        })

        return out

    def feed_forward_pure(self, x):

        out = self.sess.run(self.network, feed_dict = {

            self.input_tensor: x

        })

        return out

    def feed_forward_features(self, x):

        out = self.sess.run(self.feature_tensor, feed_dict = {

            self.input_tensor: x,

        })

        return out

    def test_euclidean_distance(self, x, points, present_joints, inside_box_joints):

        err = self.sess.run(self.euclidean_dist, feed_dict = {

            self.input_tensor: x,

            self.desired_points: points,

            self.present_joints: present_joints,

            self.inside_box_joints: inside_box_joints

        })

        return err

    def test_joint_distances(self, x, y):

        err = self.sess.run(self.euclidean_distance(), feed_dict = {

            self.input_tensor: x,

            self.desired_points: y

        })

        return err

    def test_joint_activations(self, x):

        err = self.sess.run(self._joint_highest_activations(), feed_dict = {

            self.input_tensor: x

        })

        return err

    def estimate_joints(self, x):

        out = self.sess.run(self._joint_positions(), feed_dict = {

            self.input_tensor: x

        })

        return out

    def train(self, x, heatmaps, present_joints, learning_rate, is_inside_box):

        if not self.is_train:

            raise Exception('Network is not in train mode!')

        self.sess.run(self.optimize, feed_dict = {

            self.input_tensor: x,

            self.desired_heatmap: heatmaps,

            self.present_joints: present_joints,

            self.learning_rate: learning_rate,

            self.inside_box_joints: is_inside_box

        })

    def write_test_summary(self, epoch, loss):

        loss_sum = tf.Summary()

        loss_sum.value.add(

            tag = 'Average Euclidean Distance',

            simple_value = float(loss))

        self.summary_writer.add_summary(loss_sum, epoch)

        self.summary_writer.flush()

    def write_summary(self, inp, desired_points, heatmaps, present_joints, learning_rate, is_inside_box,

                      write_frequency = 20, write_per_joint_frequency = 100):

        step = tf.train.global_step(self.sess, self.global_step)

        if step % write_frequency == 0:

            feed_dict = {

                self.input_tensor: inp,

                self.desired_points: desired_points,

                self.desired_heatmap: heatmaps,

                self.present_joints: present_joints,

                self.learning_rate: learning_rate,

                self.inside_box_joints: is_inside_box

            }

            summary, loss = self.sess.run([self.ALL_SUMMARIES, self.loss_err], feed_dict = feed_dict)

            self.summary_writer.add_summary(summary, step)

            if step % write_per_joint_frequency == 0:

                summaries = self.sess.run(self.SUMMARIES_PER_JOINT, feed_dict = feed_dict)

                for i in range(16):

                    self.summary_writer_by_points[i].add_summary(summaries[i], step)

                for i in range(16):

                    self.summary_writer_by_points[i].flush()

            self.summary_writer.flush()

    @abstractmethod

    def pre_process(self, inp):

        pass

    @abstractmethod

    def get_network(self, input_tensor, is_training):

        pass

    @abstractmethod

    def create_summary_from_weights(self):

        pass

class HumanPoseIRNetwork(HumanPoseNN):

    """

    The first part of our network that exposes as an extractor of spatial features. It s derived from

    Inception-Resnet-v2 architecture and modified for generating heatmaps - i.e. dense predictions of body joints.

    """

    FEATURES = 32

    IMAGE_SIZE = 299

    HEATMAP_SIZE = 289

    POINT_DIAMETER = 15

    SMOOTH_SIZE = 21

    def __init__(self, log_name = None, loss_type = 'SCE', is_training = False):

        super().__init__(log_name, self.HEATMAP_SIZE, self.IMAGE_SIZE, loss_type, is_training)

    def pre_process(self, inp):

        return ((inp / 255) - 0.5) * 2.0

    def get_network(self, input_tensor, is_training):

        # Load pre-trained inception-resnet model

        with slim.arg_scope(inception_resnet_v2_arg_scope(batch_norm_decay = 0.999, weight_decay = 0.0001)):

            net, end_points = inception_resnet_v2(input_tensor, is_training = is_training)

        # Adding some modification to original InceptionResnetV2 - changing scoring of AUXILIARY TOWER

        weight_decay = 0.0005

        with tf.variable_scope('NewInceptionResnetV2'):

            with tf.variable_scope('AuxiliaryScoring'):

                with slim.arg_scope([layers.convolution2d, layers.convolution2d_transpose],

                                    weights_regularizer = slim.l2_regularizer(weight_decay),

                                    biases_regularizer = slim.l2_regularizer(weight_decay),

                                    activation_fn = None):

                    tf.summary.histogram('Last_layer/activations', net, [KEY_SUMMARIES])

                    # Scoring

                    net = slim.dropout(net, 0.7, is_training = is_training, scope = 'Dropout')

                    net = layers.convolution2d(net, num_outputs = self.FEATURES, kernel_size = 1, stride = 1,

                                               scope = 'Scoring_layer')

                    feature = net

                    tf.summary.histogram('Scoring_layer/activations', net, [KEY_SUMMARIES])

                    # Upsampling

                    net = layers.convolution2d_transpose(net, num_outputs = 16, kernel_size = 17, stride = 17,

                                                         padding = 'VALID', scope = 'Upsampling_layer')

                    tf.summary.histogram('Upsampling_layer/activations', net, [KEY_SUMMARIES])

            # Smoothing layer - separable gaussian filters

            net = super()._get_gauss_smoothing_net(net, size = self.SMOOTH_SIZE, std = 1.0, kernel_sum = 0.2)

            return net, feature

    def restore(self, checkpoint_path, is_pre_trained_imagenet_checkpoint = False):

        all_vars = tf.get_collection(tf.GraphKeys.MODEL_VARIABLES, scope = 'InceptionResnetV2')

        if not is_pre_trained_imagenet_checkpoint:

            all_vars += tf.get_collection(tf.GraphKeys.MODEL_VARIABLES, scope = 'NewInceptionResnetV2/AuxiliaryScoring')

        super()._restore(checkpoint_path, all_vars)

    def save(self, checkpoint_path, name):

        all_vars = tf.get_collection(tf.GraphKeys.MODEL_VARIABLES, scope = 'InceptionResnetV2')

        all_vars += tf.get_collection(tf.GraphKeys.MODEL_VARIABLES, scope = 'NewInceptionResnetV2/AuxiliaryScoring')

        super()._save(checkpoint_path, name, all_vars)

    def create_summary_from_weights(self):

        with tf.variable_scope('NewInceptionResnetV2/AuxiliaryScoring', reuse = True):

            tf.summary.histogram('Scoring_layer/biases', tf.get_variable('Scoring_layer/biases'), [KEY_SUMMARIES])

            tf.summary.histogram('Upsampling_layer/biases', tf.get_variable('Upsampling_layer/biases'), [KEY_SUMMARIES])

            tf.summary.histogram('Scoring_layer/weights', tf.get_variable('Scoring_layer/weights'), [KEY_SUMMARIES])

            tf.summary.histogram('Upsampling_layer/weights', tf.get_variable('Upsampling_layer/weights'),

                                 [KEY_SUMMARIES])

        with tf.variable_scope('InceptionResnetV2/AuxLogits', reuse = True):

            tf.summary.histogram('Last_layer/weights', tf.get_variable('Conv2d_2a_5x5/weights'), [KEY_SUMMARIES])

            tf.summary.histogram('Last_layer/beta', tf.get_variable('Conv2d_2a_5x5/BatchNorm/beta'), [KEY_SUMMARIES])

            tf.summary.histogram('Last_layer/moving_mean', tf.get_variable('Conv2d_2a_5x5/BatchNorm/moving_mean'),

                                 [KEY_SUMMARIES])

class PartDetector(HumanPoseNN):

    """

    Architecture of Part Detector network, as was described in https://arxiv.org/abs/1609.01743

    """

    IMAGE_SIZE = 256

    HEATMAP_SIZE = 256

    POINT_DIAMETER = 11

    def __init__(self, log_name = None, init_from_checkpoint = None, loss_type = 'SCE', is_training = False):

        if init_from_checkpoint is not None:

            part_detector.init_model_variables(init_from_checkpoint, is_training)

            self.reuse = True

        else:

            self.reuse = False

        super().__init__(log_name, self.HEATMAP_SIZE, self.IMAGE_SIZE, loss_type, is_training)

    def pre_process(self, inp):

        return inp / 255

    def create_summary_from_weights(self):

        pass

    def restore(self, checkpoint_path):

        all_vars = tf.get_collection(tf.GraphKeys.VARIABLES, scope = 'HumanPoseResnet')

        all_vars += tf.get_collection(tf.GraphKeys.MODEL_VARIABLES, scope = 'NewHumanPoseResnet/Scoring')

        super()._restore(checkpoint_path, all_vars)

    def save(self, checkpoint_path, name):

        all_vars = tf.get_collection(tf.GraphKeys.MODEL_VARIABLES, scope = 'HumanPoseResnet')

        all_vars += tf.get_collection(tf.GraphKeys.MODEL_VARIABLES, scope = 'NewHumanPoseResnet/Scoring')

        super()._save(checkpoint_path, name, all_vars)

    def get_network(self, input_tensor, is_training):

        net_end, end_points = part_detector.human_pose_resnet(input_tensor, reuse = self.reuse, training = is_training)

        return net_end, end_points['features']