# --------------------------------------------------------
# Tensorflow Faster R-CNN
# Licensed under The MIT License [see LICENSE for details]
# Written by Xinlei Chen
# --------------------------------------------------------

from tflearn.layers.conv import global_avg_pool, global_max_pool
import tensorflow as tf
import tensorflow.contrib.slim as slim
import numpy as np
import lib.config.config as cfg
from lib.nets.network import Network

class vgg16(Network):
    def __init__(self, batch_size=1):
        Network.__init__(self, batch_size=batch_size)

    def build_network(self, sess, is_training=True):
        with tf.variable_scope('vgg_16', 'vgg_16'):

            # select initializer
            if cfg.FLAGS.initializer == "truncated":
                initializer = tf.truncated_normal_initializer(mean=0.0, stddev=0.01)
                initializer_bbox = tf.truncated_normal_initializer(mean=0.0, stddev=0.001)
            else:
                initializer = tf.random_normal_initializer(mean=0.0, stddev=0.01)
                initializer_bbox = tf.random_normal_initializer(mean=0.0, stddev=0.001)

            # Build head
            net = self.build_head(is_training)

            # Build rpn
            rpn_cls_prob, rpn_bbox_pred, rpn_cls_score, rpn_cls_score_reshape = self.build_rpn(net, is_training, initializer)

            # Build proposals
            rois = self.build_proposals(is_training, rpn_cls_prob, rpn_bbox_pred, rpn_cls_score)

            # Build predictions
            cls_score, cls_prob, bbox_pred = self.build_predictions(net, rois, is_training, initializer, initializer_bbox)

            self._predictions["rpn_cls_score"] = rpn_cls_score
            self._predictions["rpn_cls_score_reshape"] = rpn_cls_score_reshape
            self._predictions["rpn_cls_prob"] = rpn_cls_prob
            self._predictions["rpn_bbox_pred"] = rpn_bbox_pred
            self._predictions["cls_score"] = cls_score
            self._predictions["cls_prob"] = cls_prob
            self._predictions["bbox_pred"] = bbox_pred
            self._predictions["rois"] = rois

            self._score_summaries.update(self._predictions)

            return rois, cls_prob, bbox_pred

    def get_variables_to_restore(self, variables, var_keep_dic):
        variables_to_restore = []

        for v in variables:
            # exclude the conv weights that are fc weights in vgg16
            if v.name == 'vgg_16/fc6/weights:0' or v.name == 'vgg_16/fc7/weights:0':
                self._variables_to_fix[v.name] = v
                continue
            # exclude the first conv layer to swap RGB to BGR
            if v.name == 'vgg_16/conv1/conv1_1/weights:0':
                self._variables_to_fix[v.name] = v
                continue
            if v.name.split(':')[0] in var_keep_dic:
                print('Variables restored: %s' % v.name)
                variables_to_restore.append(v)

        return variables_to_restore

    def fix_variables(self, sess, pretrained_model):
        print('Fix VGG16 layers..')
        with tf.variable_scope('Fix_VGG16'):
            with tf.device("/cpu:0"):
                # fix the vgg16 issue from conv weights to fc weights
                # fix RGB to BGR
                fc6_conv = tf.get_variable("fc6_conv", [7, 7, 512, 4096], trainable=False)
                fc7_conv = tf.get_variable("fc7_conv", [1, 1, 4096, 4096], trainable=False)
                conv1_rgb = tf.get_variable("conv1_rgb", [3, 3, 3, 64], trainable=False)
                restorer_fc = tf.train.Saver({"vgg_16/fc6/weights": fc6_conv,
                                              "vgg_16/fc7/weights": fc7_conv,
                                              "vgg_16/conv1/conv1_1/weights": conv1_rgb})
                restorer_fc.restore(sess, pretrained_model)

                sess.run(tf.assign(self._variables_to_fix['vgg_16/fc6/weights:0'], tf.reshape(fc6_conv,
                                                                                              self._variables_to_fix['vgg_16/fc6/weights:0'].get_shape())))
                sess.run(tf.assign(self._variables_to_fix['vgg_16/fc7/weights:0'], tf.reshape(fc7_conv,
                                                                                              self._variables_to_fix['vgg_16/fc7/weights:0'].get_shape())))
                sess.run(tf.assign(self._variables_to_fix['vgg_16/conv1/conv1_1/weights:0'],
                                   tf.reverse(conv1_rgb, [2])))

    def build_head(self, is_training):

        # Main network
        # Layer  1
        net = slim.repeat(self._image, 2, slim.conv2d, 64, [3, 3], trainable=False, scope='conv1')
        net = slim.max_pool2d(net, [2, 2], padding='SAME', scope='pool1')

        # Layer 2
        net = slim.repeat(net, 2, slim.conv2d, 128, [3, 3], trainable=False, scope='conv2')
        net = slim.max_pool2d(net, [2, 2], padding='SAME', scope='pool2')
        # Layer 3
        net = slim.conv2d(net, 256, [3, 3], trainable=is_training, scope='conv3/conv3_1')
        self._layers['3_1'] = net
        net = slim.conv2d(net, 256, [3, 3], trainable=is_training, scope='conv3/conv3_2')
        self._layers['3_2'] = net
        net = slim.conv2d(net, 256, [3, 3], trainable=is_training, scope='conv3/conv3_3')
        self._layers['3_3'] = net
        net = slim.max_pool2d(net, [2, 2], padding='SAME', scope='pool3')
        # Layer 4
        net = slim.conv2d(net, 512, [3, 3], trainable=is_training, scope='conv4/conv4_1')
        self._layers['4_1'] = net
        net = slim.conv2d(net, 512, [3, 3], trainable=is_training, scope='conv4/conv4_2')
        self._layers['4_2'] = net
        net = slim.conv2d(net, 512, [3, 3], trainable=is_training, scope='conv4/conv4_3')
        self._layers['4_3'] = net
        net = slim.max_pool2d(net, [2, 2], padding='SAME', scope='pool4')
        # Layer 5
        net = slim.conv2d(net, 512, [3, 3], trainable=is_training, scope='conv5/conv5_1')
        self._layers['5_1'] = net
        net = slim.conv2d(net, 512, [3, 3], trainable=is_training, scope='conv5/conv5_2')
        self._layers['5_2'] = net
        net = slim.conv2d(net, 512, [3, 3], trainable=is_training, scope='conv5/conv5_3')
        self._layers['5_3'] = net

        # Append network to summaries
        self._act_summaries.append(net)

        # Append network as head layer
        self._layers['head'] = net

        return net

    def build_rpn(self, net, is_training, initializer):

        # Build anchor component
        self._anchor_component()

        # Create RPN Layer

        rpn = slim.conv2d(net, 512, [3, 3], trainable=is_training, weights_initializer=initializer,scope="rpn_conv/3x3")
      #  rpn_first = slim.conv2d(net, 512, [1, 1], trainable=is_training, weights_initializer=initializer, scope="rpn_conv/3x3")
       # self._act_summaries.append(rpn_first)
      #  rpn = convolutional_block_attention_module_SC(rpn_first, 1)

        #self._act_summaries.append(rpn)

        rpn_cls_score = slim.conv2d(rpn, self._num_anchors * 2, [1, 1], trainable=is_training, weights_initializer=initializer, padding='VALID', activation_fn=None, scope='rpn_cls_score')

        # Change it so that the score has 2 as its channel size
        rpn_cls_score_reshape = self._reshape_layer(rpn_cls_score, 2, 'rpn_cls_score_reshape')
        rpn_cls_prob_reshape = self._softmax_layer(rpn_cls_score_reshape, "rpn_cls_prob_reshape")
        rpn_cls_prob = self._reshape_layer(rpn_cls_prob_reshape, self._num_anchors * 2, "rpn_cls_prob")
        rpn_bbox_pred = slim.conv2d(rpn,self._num_anchors * 4, [1, 1], trainable=is_training, weights_initializer=initializer, padding='VALID', activation_fn=None, scope='rpn_bbox_pred')

        return rpn_cls_prob, rpn_bbox_pred, rpn_cls_score, rpn_cls_score_reshape

    def build_proposals(self, is_training, rpn_cls_prob, rpn_bbox_pred, rpn_cls_score):

        if is_training:
            rois, roi_scores = self._proposal_layer(rpn_cls_prob, rpn_bbox_pred, "rois")
            rpn_labels = self._anchor_target_layer(rpn_cls_score, "anchor")

            # Try to have a deterministic order for the computing graph, for reproducibility
            with tf.control_dependencies([rpn_labels]):
                rois, _ = self._proposal_target_layer(rois, roi_scores, "rpn_rois")
        else:
            if cfg.FLAGS.test_mode == 'nms':
                rois, _ = self._proposal_layer(rpn_cls_prob, rpn_bbox_pred, "rois")
            elif cfg.FLAGS.test_mode == 'top':
                rois, _ = self._proposal_top_layer(rpn_cls_prob, rpn_bbox_pred, "rois")
            else:
                raise NotImplementedError
        return rois

    def build_predictions(self, net, rois, is_training, initializer, initializer_bbox):

        # Crop image ROIs
        pool5 = self._crop_pool_layer(net, rois, "pool5")
        pool5_flat = slim.flatten(pool5, scope='flatten')

        # Fully connected layers
        fc6 = slim.fully_connected(pool5_flat, 4096, scope='fc6')
        if is_training:
            fc6 = slim.dropout(fc6, keep_prob=0.5, is_training=True, scope='dropout6')

        fc7 = slim.fully_connected(fc6, 4096, scope='fc7')
        if is_training:
            fc7 = slim.dropout(fc7, keep_prob=0.5, is_training=True, scope='dropout7')

        # Scores and predictions
        cls_score = slim.fully_connected(fc7, self._num_classes, weights_initializer=initializer, trainable=is_training, activation_fn=None, scope='cls_score')
        cls_prob = self._softmax_layer(cls_score, "cls_prob")
        bbox_prediction = slim.fully_connected(fc7, self._num_classes * 4, weights_initializer=initializer_bbox, trainable=is_training, activation_fn=None, scope='bbox_pred')

        return cls_score, cls_prob, bbox_prediction

def combined_static_and_dynamic_shape(tensor):
  """Returns a list containing static and dynamic values for the dimensions.

  Returns a list of static and dynamic values for shape dimensions. This is
  useful to preserve static shapes when available in reshape operation.

  Args:
    tensor: A tensor of any type.

  Returns:
    A list of size tensor.shape.ndims containing integers or a scalar tensor.
  """
  static_tensor_shape = tensor.shape.as_list()
  dynamic_tensor_shape = tf.shape(tensor)
  combined_shape = []
  for index, dim in enumerate(static_tensor_shape):
    if dim is not None:
      combined_shape.append(dim)
    else:
      combined_shape.append(dynamic_tensor_shape[index])
  return combined_shape

def convolutional_block_attention_module(feature_map, index, inner_units_ratio=0.5):
    """
    CBAM: convolution block attention module, which is described in "CBAM: Convolutional Block Attention Module"
    Architecture : "https://arxiv.org/pdf/1807.06521.pdf"
    If you want to use this module, just plug this module into your network
    :param feature_map : input feature map
    :param index : the index of convolution block attention module
    :param inner_units_ratio: output units number of fully connected layer: inner_units_ratio*feature_map_channel
    :return:feature map with channel and spatial attention
    """
    with tf.variable_scope("cbam_%s" % (index)):
        feature_map_shape = combined_static_and_dynamic_shape(feature_map)
        # feature_map_shape[0] = 1
        # feature_map_shape[1] = 32
        # feature_map_shape[2] = 32
        # feature_map_shape[3] = 512
        # channel attention

        # channel_avg_weights = tf.nn.avg_pool(
        #     value=feature_map,
        #     ksize=[1, feature_map_shape[1], feature_map_shape[2], 1],
        #     strides=[1, 1, 1, 1],
        #     padding='VALID'
        # )
        # channel_max_weights = tf.nn.max_pool(
        #     value=feature_map,
        #     ksize=[1, feature_map_shape[1], feature_map_shape[2], 1],
        #     strides=[1, 1, 1, 1],
        #     padding='VALID'
        # )
        globel_avg = global_avg_pool(feature_map)
        channel_avg_weights = tf.reshape(globel_avg, [1, 1, 1, 512])
        globel_max = global_max_pool(feature_map)
        channel_max_weights = tf.reshape(globel_max, [1, 1, 1, 512])

        channel_avg_reshape = tf.reshape(channel_avg_weights,
                                         [feature_map_shape[0], 1, feature_map_shape[3]])
        channel_max_reshape = tf.reshape(channel_max_weights,
                                         [feature_map_shape[0], 1, feature_map_shape[3]])
        # channel_w_reshape = tf.concat([channel_avg_reshape, channel_max_reshape], axis=1)
        #
        # fc_1 = tf.layers.dense(
        #     inputs=channel_w_reshape,
        #     units=feature_map_shape[3] * inner_units_ratio,
        #     name="fc_1",
        #     activation=tf.nn.relu
        # )
        # fc_2 = tf.layers.dense(
        #     inputs=fc_1,
        #     units=feature_map_shape[3],
        #     name="fc_2",
        #     activation=None
        # )
        channel_w_reshape = tf.concat([channel_avg_reshape, channel_max_reshape], axis=3)

        fc_1 = slim.conv2d(
            channel_w_reshape,
            feature_map_shape[3] * inner_units_ratio,
            [1, 1],
            padding='SAME',
            activation_fn=tf.nn.sigmoid,
            scope="channel_attention_conv1"
        )
        fc_2 = slim.conv2d(
            fc_1,
            feature_map_shape[3],
            [1, 1],
            padding='SAME',
            activation_fn=tf.nn.sigmoid,
            scope="channel_attention_conv2"
        )
        channel_attention = tf.reduce_sum(fc_2, axis=1, name="channel_attention_sum")
        channel_attention = tf.nn.sigmoid(channel_attention, name="channel_attention_sum_sigmoid")
        channel_attention = tf.reshape(channel_attention, shape=[feature_map_shape[0], 1, 1, feature_map_shape[3]])
        feature_map_with_channel_attention = tf.multiply(feature_map, channel_attention)
        # spatial attention
        channel_wise_avg_pooling = tf.reduce_mean(feature_map_with_channel_attention, axis=3)
        channel_wise_max_pooling = tf.reduce_max(feature_map_with_channel_attention, axis=3)

        channel_wise_avg_pooling = tf.reshape(channel_wise_avg_pooling,
                                              shape=[feature_map_shape[0], feature_map_shape[1], feature_map_shape[2], 1])
        channel_wise_max_pooling = tf.reshape(channel_wise_max_pooling,
                                              shape=[feature_map_shape[0], feature_map_shape[1], feature_map_shape[2], 1])

        channel_wise_pooling = tf.concat([channel_wise_avg_pooling, channel_wise_max_pooling], axis=3)
        spatial_attention = slim.conv2d(
            channel_wise_pooling,
            1,
            [7, 7],
            padding='SAME',
            activation_fn=tf.nn.sigmoid,
            scope="spatial_attention_conv"
        )
        feature_map_with_attention = tf.multiply(feature_map_with_channel_attention, spatial_attention)
        return feature_map_with_attention
        # return feature_map_with_channel_attention
def convolutional_block_attention_module_SC(feature_map, index, inner_units_ratio=0.5):
    """
    CBAM: convolution block attention module, which is described in "CBAM: Convolutional Block Attention Module"
    Architecture : "https://arxiv.org/pdf/1807.06521.pdf"
    If you want to use this module, just plug this module into your network
    :param feature_map : input feature map
    :param index : the index of convolution block attention module
    :param inner_units_ratio: output units number of fully connected layer: inner_units_ratio*feature_map_channel
    :return:feature map with channel and spatial attention
    """
    with tf.variable_scope("cbam_%s" % (index)):
        feature_map_shape = combined_static_and_dynamic_shape(feature_map)

        # //spatial attention//
        channel_wise_avg_pooling = tf.reduce_mean(feature_map, axis=3)
        channel_wise_max_pooling = tf.reduce_max(feature_map, axis=3)

        channel_wise_avg_pooling = tf.reshape(channel_wise_avg_pooling,
                                              shape=[feature_map_shape[0], feature_map_shape[1], feature_map_shape[2],
                                                     1])
        channel_wise_max_pooling = tf.reshape(channel_wise_max_pooling,
                                              shape=[feature_map_shape[0], feature_map_shape[1], feature_map_shape[2],
                                                     1])

        channel_wise_pooling = tf.concat([channel_wise_avg_pooling, channel_wise_max_pooling], axis=3)
        spatial_attention = slim.conv2d(
            channel_wise_pooling,
            1,
            [7, 7],
            padding='SAME',
            activation_fn=tf.nn.sigmoid,
            scope="spatial_attention_conv"
        )
        feature_map_with_spatial_attention = tf.multiply(feature_map, spatial_attention)
        # //channel attention//
        # channel_avg_weights = tf.nn.avg_pool(
        #     value=feature_map,
        #     ksize=[1, feature_map_shape[1], feature_map_shape[2], 1],
        #     strides=[1, 1, 1, 1],
        #     padding='VALID'
        # )
        # channel_max_weights = tf.nn.max_pool(
        #     value=feature_map,
        #     ksize=[1, feature_map_shape[1], feature_map_shape[2], 1],
        #     strides=[1, 1, 1, 1],
        #     padding='VALID'
        # )
        globel_avg = global_avg_pool(feature_map_with_spatial_attention)
        channel_avg_weights = tf.reshape(globel_avg, [1, 1, 1, 512])
        globel_max = global_max_pool(feature_map_with_spatial_attention)
        channel_max_weights = tf.reshape(globel_max, [1, 1, 1, 512])

        # //original program//
        # channel_avg_reshape = tf.reshape(channel_avg_weights,
        #                                  [feature_map_shape[0], 1, feature_map_shape[3]])
        # channel_max_reshape = tf.reshape(channel_max_weights,
        #                                  [feature_map_shape[0], 1, feature_map_shape[3]])
        # channel_w_reshape = tf.concat([channel_avg_reshape, channel_max_reshape], axis=1)
        #
        # fc_1 = tf.layers.dense(
        #     inputs=channel_w_reshape,
        #     units=feature_map_shape[3] * inner_units_ratio,
        #     name="fc_1",
        #     activation=tf.nn.relu
        # )
        # fc_2 = tf.layers.dense(
        #     inputs=fc_1,
        #     units=feature_map_shape[3],
        #     name="fc_2",
        #     activation=None
        # )

        # //improved program//
        channel_avg_reshape = tf.reshape(channel_avg_weights,
                                         [feature_map_shape[0], 1, 1, feature_map_shape[3]])
        channel_max_reshape = tf.reshape(channel_max_weights,
                                         [feature_map_shape[0], 1, 1, feature_map_shape[3]])
        channel_w_reshape = tf.concat([channel_avg_reshape, channel_max_reshape], axis=3)

        fc_1 = slim.conv2d(
            channel_w_reshape,
            feature_map_shape[3] * inner_units_ratio,
            [1, 1],
            padding='SAME',
            activation_fn=tf.nn.sigmoid,
            scope="channel_attention_conv1"
        )
        fc_2 = slim.conv2d(
            fc_1,
            feature_map_shape[3],
            [1, 1],
            padding='SAME',
            activation_fn=tf.nn.sigmoid,
            scope="channel_attention_conv2"
        )

        channel_attention = tf.reduce_sum(fc_2, axis=1, name="channel_attention_sum")
        channel_attention = tf.nn.sigmoid(channel_attention, name="channel_attention_sum_sigmoid")
        channel_attention = tf.reshape(channel_attention, shape=[feature_map_shape[0], 1, 1, feature_map_shape[3]])
        feature_map_with_channel_attention = tf.multiply(feature_map_with_spatial_attention, channel_attention)

        return feature_map_with_channel_attention