from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
import tf_slim as slim
# from tensorflow.contrib import slim
from models.specs import TFNetwork


def tf_nms_filter(pred_boxes, scores_vector, nms_overlap_thresh, K, L):
    """
    :param pred_boxes: Tensor(float32) (num_boxes, 4)
    :param scores_vector: Tensor(float32) (num_boxes, )
    :param nms_overlap_thresh: (float)
    :param K: (int) maximum number of boxes after NMS
    :param L: (int) number of classes to sum for score calculation
    :return:
    """
    # get predicted boxes and their scores
    scores = tf.reduce_sum(scores_vector[:, -L:], axis=1)
    idx = tf.image.non_max_suppression(pred_boxes, scores, iou_threshold=nms_overlap_thresh, max_output_size=K,
                                       name='good_boxes_idx')
    return idx


def get_full_feature_maps(heatmap_features, regression_features, name):
    box_filter_input = tf.concat([heatmap_features, regression_features], -1, name=name)
    return box_filter_input


def iou_prediction(features, num_classes, scope, reuse=None, L2_reg=0.0, act_func=tf.nn.relu):
    """ Transform features to IoU prediction"""

    def _args_scope():
        with slim.arg_scope([slim.conv2d], activation_fn=act_func, weights_regularizer=slim.l2_regularizer(L2_reg)):
            with slim.arg_scope([slim.conv2d], padding='SAME') as arg_sc:
                return arg_sc

    with slim.arg_scope(_args_scope()):
        with tf.compat.v1.variable_scope(scope, scope, [features], reuse=reuse) as sc:
            end_points_collection = sc.name + '_end_points'
            # Collect outputs for conv2d, fully_connected and max_pool2d.
            with slim.arg_scope([slim.conv2d, slim.fully_connected, slim.max_pool2d],
                                outputs_collections=end_points_collection):
                logits = slim.conv2d(features, num_classes, [1, 1], stride=1, activation_fn=None,
                                     scope='box_logit_conv')
                logits = tf.squeeze(logits, axis=[1, 2], name='box_logit')

    return logits


def iou_loss(y, y_hat, scope, reduction=tf.reduce_mean):
    with tf.compat.v1.variable_scope('%s/loss' % scope, values=[y, y_hat]):
        cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y, logits=y_hat, name='cross_entropy')
        mean_xent = reduction(cross_entropy, name='mean_xent_loss')
    return mean_xent


def phoc_loss_func(y, y_hat, scope, reduction=tf.reduce_mean):
    """
    Take phoc assignments to box proposals (y) and calculate loss relative to predicted
    PHOCs (y_hat)
    :param y: (tf.Tensor) assigned PHOCs e.g. by random_rois op
    :param y_hat: predicted PHOCs
    :param y_assigned_labels: proposal box labels also generated by random_rois op (self.random_boxes_labels)
    :param iou_lables_thresh: (int) what IoU coverage classes should we take for PHOC training.
    :return: (scalar)
    """
    # gy = tf.squeeze(tf.gather(y, tf.where(y_assigned_labels >= iou_lables_thresh)), name='gt_phocs_for_loss')
    # gy_hat = tf.squeeze(tf.gather(y_hat, tf.where(y_assigned_labels >= iou_lables_thresh)), name='pred_phocs_for_loss')
    with tf.compat.v1.variable_scope('%s/loss' % scope, values=[y, y_hat]):
        xent_vec = tf.nn.sigmoid_cross_entropy_with_logits(labels=y, logits=y_hat, name='phoc_sigmoid_xent_vec')
        loss = reduction(xent_vec, name='phoc_loss')
    return loss


def _args_scope(act_func, L2_reg):
    with slim.arg_scope([slim.conv2d],
                        activation_fn=act_func,
                        weights_regularizer=slim.l2_regularizer(L2_reg)
                        ):
        with slim.arg_scope([slim.conv2d], padding='SAME') as arg_sc:
            return arg_sc


class IoUPrediction(TFNetwork):

    def __init__(self, output_shape, max_boxes_to_filter=600, act_func=tf.nn.relu, **kwargs):
        super(IoUPrediction, self).__init__(**kwargs)

        self.act_func = act_func
        self.output_shape = output_shape
        self._roi_pool_call = 0
        self._iou_calls = 0

        # Filtering bounding boxes
        self.max_boxes_to_filter = max_boxes_to_filter
        self.good_boxes_idx = None
        self.good_idx = None

    def _pooling(self, x, b, func):
        """Prepare boxes for pooling and apply pooling network (defined by func)"""
        reuse = self.get_reuse(self._roi_pool_call)
        self._roi_pool_call += 1
        scope = self.scope
        L2_reg = self.args.box_filter_L2_reg
        dropout = self.args.dropout

        def _args_scope():
            with slim.arg_scope([slim.conv2d, slim.fully_connected], activation_fn=self.act_func,
                                weights_regularizer=slim.l2_regularizer(L2_reg)
                                ):
                with slim.arg_scope([slim.conv2d], padding='SAME') as arg_sc:
                    return arg_sc

        with slim.arg_scope(_args_scope()):
            with tf.compat.v1.variable_scope(scope, scope, [x, b], reuse=reuse) as sc:
                end_points_collection = sc.name + '_end_points'
                # Collect outputs for conv2d, fully_connected and max_pool2d.
                with slim.arg_scope([slim.conv2d, slim.fully_connected, slim.max_pool2d],
                                    outputs_collections=end_points_collection):
                    boxes_input = tf.identity(b[:, 1:], name='boxes')
                    batch_idx = tf.cast(b[:, 0], dtype=tf.int32, name='batch_idx')
                    net = func(x, boxes_input, batch_idx)
        return net

    def base_pooling(self, x, b):
        def base(x, boxes_input, batch_idx):
            pooled_features = tf.image.crop_and_resize(x, boxes_input, batch_idx, crop_size=self.output_shape)
            net = slim.conv2d(pooled_features, 64, [5, 5], padding='VALID', stride=2, scope='conv1_box_cls')
            net = slim.conv2d(net, 128, [3, 3], padding='VALID', stride=2, scope='conv2_box_cls')
            net = slim.conv2d(net, 256, [3, 3], padding='SAME', stride=1, scope='conv3_box_cls')
            net = slim.conv2d(net, 512, [3, 3], padding='SAME', stride=1, scope='conv4_box_cls')
            net = slim.conv2d(net, 512, [2, 14], padding='VALID', stride=1, scope='conv5_box_cls')
            net = slim.conv2d(net, 512, [1, 1], padding='SAME', stride=1, scope='conv6_box_cls')
            net = slim.conv2d(net, 1024, [1, 1], padding='SAME', stride=1, scope='conv7_box_cls')

            conv_comp_boxes = tf.reshape(boxes_input, [-1, 1, 1, 4])
            box_context = slim.conv2d(conv_comp_boxes, 8, [1, 1], stride=1, scope='fc1_context')
            box_context = slim.conv2d(box_context, 32, [1, 1], stride=1, scope='fc2_context')
            box_context = slim.conv2d(box_context, 128, [1, 1], stride=1, scope='fc3_context')
            net = tf.concat([net, box_context], axis=-1, name='feature_with_context')
            return net

        net = self._pooling(x, b, base)
        return net

    def large_pooling(self, x, b):
        def large(x, boxes_input, batch_idx):
            pooled_features = tf.image.crop_and_resize(x, boxes_input, batch_idx, crop_size=self.output_shape)
            net = slim.conv2d(pooled_features, 1024, self.output_shape, padding='VALID', stride=1,
                              scope='conv1_box_cls')
            net = slim.conv2d(net, 1024, [1, 1], stride=1, scope='conv2_box_cls')
            net = slim.conv2d(net, 1024, [1, 1], stride=1, padding='VALID', scope='fc1_box_cls')
            net = slim.conv2d(net, 1024, [1, 1], stride=1, padding='VALID', scope='fc2_box_cls')
            net = slim.conv2d(net, 1024, [1, 1], stride=1, padding='VALID', scope='fc3_box_cls')

            conv_comp_boxes = tf.reshape(boxes_input, [-1, 1, 1, 4])
            box_context = slim.conv2d(conv_comp_boxes, 8, [1, 1], stride=1, scope='fc1_context')
            box_context = slim.conv2d(box_context, 32, [1, 1], stride=1, scope='fc2_context')
            box_context = slim.conv2d(box_context, 128, [1, 1], stride=1, scope='fc3_context')
            net = tf.concat([net, box_context], axis=-1, name='feature_with_context')
            return net

        net = self._pooling(x, b, large)
        return net

    def iou(self, features):
        reuse = self.get_reuse(self._iou_calls)
        self._iou_calls += 1
        iou_pred = iou_prediction(features, num_classes=self.args.box_filter_num_clsses, scope=self.scope,
                                  L2_reg=self.args.box_filter_L2_reg, reuse=reuse)
        return iou_pred

    def get_good_boxes(self, logits, pool_boxes):
        """
        Filter boxes:
            * perfom nms on boxes
            * after nms keep boxes with score above certain threshold
        """
        with tf.name_scope(self.scope):
            pred_iou_prob = tf.compat.v1.nn.softmax(logits, dim=-1, name='pred_iou_prob')
            good_boxes_idx = tf_nms_filter(pool_boxes[:, 1:], pred_iou_prob,
                                           nms_overlap_thresh=self.args.nms_overlap_thresh,
                                           L=self.args.num_of_classes_to_sum,
                                           K=self.max_boxes_to_filter)
            intermed_probs = tf.gather(pred_iou_prob, good_boxes_idx, name='intermed_probs')
            intermed_boxes = tf.gather(pool_boxes, good_boxes_idx, name='intermed_boxes')
            intermed_score = tf.reduce_sum(intermed_probs[:, -self.args.num_of_classes_to_sum:], axis=1)
            good_idx = tf.reshape(tf.where(intermed_score > self.args.score_thresh), [-1])
            good_boxes = tf.gather(intermed_boxes, good_idx, name='good_boxes')

            self.good_boxes_idx = good_boxes_idx
            self.good_idx = good_idx

        return good_boxes

    def get_good_embedding(self, embeddings):
        """
        In case we have word embeddings, filter them to fit those boxes we pick as good boxes from self.get_good_boxes
        """
        intermed_phocs = tf.gather(embeddings, self.good_boxes_idx, name='intermed_phocs')
        good_phocs = tf.gather(intermed_phocs, self.good_idx, name='good_phocs')
        return good_phocs

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
import tf_slim as slim
# from tensorflow.contrib import slim
from models.specs import TFNetwork


def tf_nms_filter(pred_boxes, scores_vector, nms_overlap_thresh, K, L):
    """
    :param pred_boxes: Tensor(float32) (num_boxes, 4)
    :param scores_vector: Tensor(float32) (num_boxes, )
    :param nms_overlap_thresh: (float)
    :param K: (int) maximum number of boxes after NMS
    :param L: (int) number of classes to sum for score calculation
    :return:
    """
    # get predicted boxes and their scores
    scores = tf.reduce_sum(scores_vector[:, -L:], axis=1)
    idx = tf.image.non_max_suppression(pred_boxes, scores, iou_threshold=nms_overlap_thresh, max_output_size=K,
                                       name='good_boxes_idx')
    return idx


def get_full_feature_maps(heatmap_features, regression_features, name):
    box_filter_input = tf.concat([heatmap_features, regression_features], -1, name=name)
    return box_filter_input


def iou_prediction(features, num_classes, scope, reuse=None, L2_reg=0.0, act_func=tf.nn.relu):
    """ Transform features to IoU prediction"""

    def _args_scope():
        with slim.arg_scope([slim.conv2d], activation_fn=act_func, weights_regularizer=slim.l2_regularizer(L2_reg)):
            with slim.arg_scope([slim.conv2d], padding='SAME') as arg_sc:
                return arg_sc

    with slim.arg_scope(_args_scope()):
        with tf.compat.v1.variable_scope(scope, scope, [features], reuse=reuse) as sc:
            end_points_collection = sc.name + '_end_points'
            # Collect outputs for conv2d, fully_connected and max_pool2d.
            with slim.arg_scope([slim.conv2d, slim.fully_connected, slim.max_pool2d],
                                outputs_collections=end_points_collection):
                logits = slim.conv2d(features, num_classes, [1, 1], stride=1, activation_fn=None,
                                     scope='box_logit_conv')
                logits = tf.squeeze(logits, axis=[1, 2], name='box_logit')

    return logits


def iou_loss(y, y_hat, scope, reduction=tf.reduce_mean):
    with tf.compat.v1.variable_scope('%s/loss' % scope, values=[y, y_hat]):
        cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y, logits=y_hat, name='cross_entropy')
        mean_xent = reduction(cross_entropy, name='mean_xent_loss')
    return mean_xent


def phoc_loss_func(y, y_hat, scope, reduction=tf.reduce_mean):
    """
    Take phoc assignments to box proposals (y) and calculate loss relative to predicted
    PHOCs (y_hat)
    :param y: (tf.Tensor) assigned PHOCs e.g. by random_rois op
    :param y_hat: predicted PHOCs
    :param y_assigned_labels: proposal box labels also generated by random_rois op (self.random_boxes_labels)
    :param iou_lables_thresh: (int) what IoU coverage classes should we take for PHOC training.
    :return: (scalar)
    """
    # gy = tf.squeeze(tf.gather(y, tf.where(y_assigned_labels >= iou_lables_thresh)), name='gt_phocs_for_loss')
    # gy_hat = tf.squeeze(tf.gather(y_hat, tf.where(y_assigned_labels >= iou_lables_thresh)), name='pred_phocs_for_loss')
    with tf.compat.v1.variable_scope('%s/loss' % scope, values=[y, y_hat]):
        xent_vec = tf.nn.sigmoid_cross_entropy_with_logits(labels=y, logits=y_hat, name='phoc_sigmoid_xent_vec')
        loss = reduction(xent_vec, name='phoc_loss')
    return loss


def _args_scope(act_func, L2_reg):
    with slim.arg_scope([slim.conv2d],
                        activation_fn=act_func,
                        weights_regularizer=slim.l2_regularizer(L2_reg)
                        ):
        with slim.arg_scope([slim.conv2d], padding='SAME') as arg_sc:
            return arg_sc


class IoUPrediction(TFNetwork):

    def __init__(self, output_shape, max_boxes_to_filter=600, act_func=tf.nn.relu, **kwargs):
        super(IoUPrediction, self).__init__(**kwargs)

        self.act_func = act_func
        self.output_shape = output_shape
        self._roi_pool_call = 0
        self._iou_calls = 0

        # Filtering bounding boxes
        self.max_boxes_to_filter = max_boxes_to_filter
        self.good_boxes_idx = None
        self.good_idx = None

    def _pooling(self, x, b, func):
        """Prepare boxes for pooling and apply pooling network (defined by func)"""
        reuse = self.get_reuse(self._roi_pool_call)
        self._roi_pool_call += 1
        scope = self.scope
        L2_reg = self.args.box_filter_L2_reg
        dropout = self.args.dropout

        def _args_scope():
            with slim.arg_scope([slim.conv2d, slim.fully_connected], activation_fn=self.act_func,
                                weights_regularizer=slim.l2_regularizer(L2_reg)
                                ):
                with slim.arg_scope([slim.conv2d], padding='SAME') as arg_sc:
                    return arg_sc

        with slim.arg_scope(_args_scope()):
            with tf.compat.v1.variable_scope(scope, scope, [x, b], reuse=reuse) as sc:
                end_points_collection = sc.name + '_end_points'
                # Collect outputs for conv2d, fully_connected and max_pool2d.
                with slim.arg_scope([slim.conv2d, slim.fully_connected, slim.max_pool2d],
                                    outputs_collections=end_points_collection):
                    boxes_input = tf.identity(b[:, 1:], name='boxes')
                    batch_idx = tf.cast(b[:, 0], dtype=tf.int32, name='batch_idx')
                    net = func(x, boxes_input, batch_idx)
        return net

    def base_pooling(self, x, b):
        def base(x, boxes_input, batch_idx):
            pooled_features = tf.image.crop_and_resize(x, boxes_input, batch_idx, crop_size=self.output_shape)
            net = slim.conv2d(pooled_features, 64, [5, 5], padding='VALID', stride=2, scope='conv1_box_cls')
            net = slim.conv2d(net, 128, [3, 3], padding='VALID', stride=2, scope='conv2_box_cls')
            net = slim.conv2d(net, 256, [3, 3], padding='SAME', stride=1, scope='conv3_box_cls')
            net = slim.conv2d(net, 512, [3, 3], padding='SAME', stride=1, scope='conv4_box_cls')
            net = slim.conv2d(net, 512, [2, 14], padding='VALID', stride=1, scope='conv5_box_cls')
            net = slim.conv2d(net, 512, [1, 1], padding='SAME', stride=1, scope='conv6_box_cls')
            net = slim.conv2d(net, 1024, [1, 1], padding='SAME', stride=1, scope='conv7_box_cls')

            conv_comp_boxes = tf.reshape(boxes_input, [-1, 1, 1, 4])
            box_context = slim.conv2d(conv_comp_boxes, 8, [1, 1], stride=1, scope='fc1_context')
            box_context = slim.conv2d(box_context, 32, [1, 1], stride=1, scope='fc2_context')
            box_context = slim.conv2d(box_context, 128, [1, 1], stride=1, scope='fc3_context')
            net = tf.concat([net, box_context], axis=-1, name='feature_with_context')
            return net

        net = self._pooling(x, b, base)
        return net

    def large_pooling(self, x, b):
        def large(x, boxes_input, batch_idx):
            pooled_features = tf.image.crop_and_resize(x, boxes_input, batch_idx, crop_size=self.output_shape)
            net = slim.conv2d(pooled_features, 1024, self.output_shape, padding='VALID', stride=1,
                              scope='conv1_box_cls')
            net = slim.conv2d(net, 1024, [1, 1], stride=1, scope='conv2_box_cls')
            net = slim.conv2d(net, 1024, [1, 1], stride=1, padding='VALID', scope='fc1_box_cls')
            net = slim.conv2d(net, 1024, [1, 1], stride=1, padding='VALID', scope='fc2_box_cls')
            net = slim.conv2d(net, 1024, [1, 1], stride=1, padding='VALID', scope='fc3_box_cls')

            conv_comp_boxes = tf.reshape(boxes_input, [-1, 1, 1, 4])
            box_context = slim.conv2d(conv_comp_boxes, 8, [1, 1], stride=1, scope='fc1_context')
            box_context = slim.conv2d(box_context, 32, [1, 1], stride=1, scope='fc2_context')
            box_context = slim.conv2d(box_context, 128, [1, 1], stride=1, scope='fc3_context')
            net = tf.concat([net, box_context], axis=-1, name='feature_with_context')
            return net

        net = self._pooling(x, b, large)
        return net

    def iou(self, features):
        reuse = self.get_reuse(self._iou_calls)
        self._iou_calls += 1
        iou_pred = iou_prediction(features, num_classes=self.args.box_filter_num_clsses, scope=self.scope,
                                  L2_reg=self.args.box_filter_L2_reg, reuse=reuse)
        return iou_pred

    def get_good_boxes(self, logits, pool_boxes):
        """
        Filter boxes:
            * perfom nms on boxes
            * after nms keep boxes with score above certain threshold
        """
        with tf.name_scope(self.scope):
            pred_iou_prob = tf.compat.v1.nn.softmax(logits, dim=-1, name='pred_iou_prob')
            good_boxes_idx = tf_nms_filter(pool_boxes[:, 1:], pred_iou_prob,
                                           nms_overlap_thresh=self.args.nms_overlap_thresh,
                                           L=self.args.num_of_classes_to_sum,
                                           K=self.max_boxes_to_filter)
            intermed_probs = tf.gather(pred_iou_prob, good_boxes_idx, name='intermed_probs')
            intermed_boxes = tf.gather(pool_boxes, good_boxes_idx, name='intermed_boxes')
            intermed_score = tf.reduce_sum(intermed_probs[:, -self.args.num_of_classes_to_sum:], axis=1)
            good_idx = tf.reshape(tf.where(intermed_score > self.args.score_thresh), [-1])
            good_boxes = tf.gather(intermed_boxes, good_idx, name='good_boxes')

            self.good_boxes_idx = good_boxes_idx
            self.good_idx = good_idx

        return good_boxes

    def get_good_embedding(self, embeddings):
        """
        In case we have word embeddings, filter them to fit those boxes we pick as good boxes from self.get_good_boxes
        """
        intermed_phocs = tf.gather(embeddings, self.good_boxes_idx, name='intermed_phocs')
        good_phocs = tf.gather(intermed_phocs, self.good_idx, name='good_phocs')
        return good_phocs

# from tensorflow.contrib import slim :param pred_boxes: Tensor(float32) (num_boxes, 4)     :param scores_vector: Tensor(float32) (num_boxes, )     :param nms_overlap_thresh: (float)     :param K: (int) maximum number of boxes after NMS     :param L: (int) number of classes to sum for score calculation     :return: # get predicted boxes and their scores Transform features to IoU prediction # Collect outputs for conv2d, fully_connected and max_pool2d. Take phoc assignments to box proposals (y) and calculate loss relative to predicted     PHOCs (y_hat)     :param y: (tf.Tensor) assigned PHOCs e.g. by random_rois op     :param y_hat: predicted PHOCs     :param y_assigned_labels: proposal box labels also generated by random_rois op (self.random_boxes_labels)     :param iou_lables_thresh: (int) what IoU coverage classes should we take for PHOC training.     :return: (scalar) # gy = tf.squeeze(tf.gather(y, tf.where(y_assigned_labels >= iou_lables_thresh)), name='gt_phocs_for_loss') # gy_hat = tf.squeeze(tf.gather(y_hat, tf.where(y_assigned_labels >= iou_lables_thresh)), name='pred_phocs_for_loss') # Filtering bounding boxes Prepare boxes for pooling and apply pooling network (defined by func) # Collect outputs for conv2d, fully_connected and max_pool2d. Filter boxes:             * perfom nms on boxes             * after nms keep boxes with score above certain threshold In case we have word embeddings, filter them to fit those boxes we pick as good boxes from self.get_good_boxes