from network.backbone import ResBlock
from network.utils import update_bboxes
import tensorflow as tf

class ROIAlign(tf.keras.Model):
    """Generate ROI features by pooling a feature map to a specified size.

    Parameters
    ----------
    pool_shape : list of ints
        Shape ([height, width]) of the ROI features generated by pooling. Defaults
        to [7, 7].
    """
    def __init__(self, pool_shape=[7, 7]):
        super().__init__()
        self.pool_shape = pool_shape

    def call(self, input_data):
        """
        Parameters
        ----------
        input_data : list of tensors
            - Feature map output by the backbone network, [N, height, width, channels]
            - ROI produced by the RPN + proposal layers; may contain zero ROIS if
              we had to pad to num_rois, [N, num_rois, (y1, x1, y2, x2)] in 
              normalized coordinates
        Returns
        -------
        tensor
            Pooled ROI, [N, num_rois, height, width, channels], height and width
            specified by pool_shape.
        """
        feature_map = input_data[0]
        roi = input_data[1]

        """Unlike the Mask R-CNN paper, we only interpolate a single value from 
        the feature map for each ROI feature pixel, rather than sampling four and 
        pooling. According to the paper this shouldn't matter to much."""
        def roi_align(x):
            features = x[0]
            boxes = x[1]
            features = tf.expand_dims(features, 0)
            inds = tf.zeros(tf.shape(boxes)[0], tf.int32)
            return tf.image.crop_and_resize(features, boxes, inds, self.pool_shape,
                                            method='bilinear')

        roi_features = tf.map_fn(roi_align, [feature_map, roi], dtype=tf.float32)

        return roi_features

class ROIHead(tf.keras.layers.Layer):
    """Classify ROI feature vectors and provide bounding box refinements for each
    class.

    For now, this is hard coded to apply the last stage of ResNet-50, but with no
    strided downsampling, then average pooling and a single hidden layer prior to
    the class/bounding-box layers, in line with the basic implementation in the
    Faster R-CNN/Mask R-CNN papers.

    Parameters
    ----------
    data_format : string
        'channels_first' or 'channels_last', indicating the ordering of feature
        maps and channels.
    pool_shape : list of ints
        Shape ([height, width]) of the ROI features generated by pooling. Defaults
        to [7, 7].
    num_channels : int
        Number of channels after convolutional stage. Defaults to 2048.
    num_classes : int
        Number of classes. Defaults to 21. Should include extra class for background.
    hidden_dim : int
        Dimension of hidden layers. Defaults to 2048.
    regularizer : function
        Regularizer function applied to all weights in the network. Defaults to
        None.
    """
    def __init__(self, data_format, pool_shape=[7, 7], num_channels=2048, 
                 num_classes=21, hidden_dim=2048, regularizer=None):
        super().__init__()
        self.res_block_1 = tf.keras.layers.TimeDistributed(ResBlock(
            num_channels, data_format, shortcut_activate=True, 
            regularizer=regularizer))
        self.res_block_2 = tf.keras.layers.TimeDistributed(ResBlock(
            num_channels, data_format, regularizer=regularizer))
        self.res_block_3 = tf.keras.layers.TimeDistributed(ResBlock(
            num_channels, data_format, regularizer=regularizer))

        bn_axis = 1 if data_format is 'channels_first' else 3
        self.squeeze_axes = [3, 4] if data_format is 'channels_first' else [2, 3]
        self.bn = tf.keras.layers.TimeDistributed(tf.layers.BatchNormalization(
            axis=bn_axis, name='bn_final'))
        self.pool = tf.keras.layers.TimeDistributed(tf.keras.layers.AveragePooling2D(
            pool_size=pool_shape, data_format=data_format))
        self.fc = tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(
            hidden_dim, name='fc', kernel_regularizer=regularizer))

        self.fc_class = tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(
            num_classes, name='fc_class', kernel_regularizer=regularizer))
        self.fc_bbox = tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(
            4*num_classes, name='fc_bbox', kernel_regularizer=regularizer))

    def call(self, input_data, training=False):
        """
        Parameters
        ----------
        input_data : tensor
            ROI feature maps produced by the ROIAlign layer, 
            [N, num_rois, height, width, channels]
        Returns
        -------
        list of tensors
            logits, probabilities, and bounding box refinements for each ROI, 
            dimensions [N, num_roi, num_classes], [N, num_roi, num_classes], 
            [N, num_roi, num_classes, (dy, dx, log(dh), log(dw))], respectively.
        """
        x = self.res_block_1(input_data, training=training)
        x = self.res_block_2(x, training=training)
        x = self.res_block_3(x, training=training)

        x = self.bn(x, training=training)
        x = tf.nn.relu(x)
        x = self.pool(x)
        x = tf.squeeze(x, axis=self.squeeze_axes)
        x = self.fc(x)
        x = tf.nn.relu(x)

        logits = self.fc_class(x)
        probs =  tf.nn.softmax(logits)
        bbox = self.fc_bbox(x)
        bbox = tf.reshape(bbox, [tf.shape(bbox)[0], tf.shape(bbox)[1], -1, 4])
        return [logits, probs, bbox]

class DetectionLayer(tf.keras.Model):
    """Generate detection events from classified ROI. Updates bounding boxes and
    filters out background boxes and low-confidence boxes.

    Parameters
    ----------
    num_detect : int
        Maximum number of objects to detect per image. Default to 100.
    prob_thresh : float
        Threshold probability (confidence) to consider a non-background ROI a 
        true object. Should be between 0 and 1. Defaults to 0.7
    overlap_thresh : float
        Threshold for deciding if proposals overlap (with respect to the IoU
        metric), during the non-max suppression stage. Defaults to 0.3.
    """
    def __init__(self, num_detect=100, prob_thresh=0.7, overlap_thresh=0.3):
        super().__init__()
        self.num_detect = num_detect        
        self.prob_thresh = prob_thresh
        self.overlap_thresh = overlap_thresh

    def call(self, input_data):
        """
        Parameters
        ----------
        input_data : list of tensors
            - ROI proposals in normalized coordinates, [N, num_rois, (y1, x1, y2, x2)]
            - Class scores of each ROI, [N, num_rois, num_classes]
            - Bounding box refinements, 
              [N, num_rois, num_classes, (dy, dx, log(dh), log(dw))]
        Returns
        -------
        tensor
            Detection boxes, in normalized coordinates, 
            [N, num_detect, (y1, x1, y2, x2, class_id, class_score)]
        """
        roi, probs, deltas = input_data
        detections = tf.map_fn(self.filter_detections, 
                               [roi, probs, deltas], dtype=tf.float32)
        return detections

    def filter_detections(self, input_data):
        """Apply bounding box deltas, filter bad boxes, and apply class-specific
        non-max suppression. This is mapped separately over each element of the 
        batch since filtering+NMS may result in different numbers of detections 
        for each batch element.

        Parameters
        ----------
        input_data : list of tensors
            - ROI proposals in normalized coordinates, [num_rois, (y1, x1, y2, x2)]
            - Class scores of each ROI, [num_rois, num_classes]
            - Bounding box refinements, 
              [num_rois, num_classes, (dy, dx, log(dh), log(dw))]
        Returns
        -------
        tensor
            Detection boxes, in normalized coordinates, 
            [num_detect, (y1, x1, y2, x2, class_id, class_score)]
        """
        roi, probs, deltas = input_data
        class_ids = tf.argmax(probs, axis=1, output_type=tf.int32)
        class_inds = tf.stack([tf.range(tf.shape(probs)[0]), class_ids], axis=1)

        class_probs = tf.gather_nd(probs, class_inds)
        class_deltas = tf.gather_nd(deltas, class_inds)
        class_roi = update_bboxes(roi, class_deltas)
        class_roi = tf.clip_by_value(class_roi, 0, 1)

        # Discard all background and low-confidence boxes
        keep = tf.where((class_ids>0) & (class_probs>=self.prob_thresh))[:, 0]
        keep = tf.cast(keep, tf.int32)
        class_ids = tf.gather(class_ids, keep)
        class_probs = tf.gather(class_probs, keep)
        class_roi = tf.gather(class_roi, keep)
        unique_class_ids = tf.unique(class_ids)[0]

        def nms(curr_id):
            indices = tf.where(tf.equal(class_ids, curr_id))[:, 0]
            curr_roi = tf.gather(class_roi, indices)
            curr_probs = tf.gather(class_probs, indices)
            class_keep = tf.image.non_max_suppression(
                curr_roi, curr_probs, self.num_detect, 
                iou_threshold=self.overlap_thresh)
            class_keep = tf.gather(indices, class_keep)
            class_keep = tf.cast(class_keep, tf.int32)
            
            # Pad with -1 so all classes return same size for stacking
            padding = tf.maximum(self.num_detect-tf.shape(class_keep)[0], 0)
            class_keep = tf.pad(class_keep, [(0, padding)], 'CONSTANT',
                               constant_values=-1)
            class_keep.set_shape([self.num_detect])
            return class_keep

        # Filter ROI for each class with NMS, up to num_detect per class, and put
        # in a single list, filtering out -1 padding we inserted in the mapping
        nms_keep = tf.map_fn(nms, unique_class_ids, dtype=tf.int32)
        # if unique_class_ids is empty, nms_keep would default to float
        nms_keep = tf.cast(nms_keep, tf.int32)
        nms_keep = tf.reshape(nms_keep, [-1])
        nms_keep = tf.gather(nms_keep, tf.where(nms_keep>-1)[:, 0])

        # Keep only top detections, up to num_detect total
        nms_class_probs = tf.gather(class_probs, nms_keep)
        num_keep = tf.minimum(tf.shape(nms_class_probs)[0], self.num_detect)
        top_probs = tf.nn.top_k(nms_class_probs, k=num_keep, sorted=True)[1]
        final_keep = tf.gather(nms_keep, top_probs)

        final_roi = tf.gather(class_roi, final_keep)
        final_ids = tf.cast(tf.gather(class_ids, final_keep), tf.float32)
        final_probs = tf.gather(class_probs, final_keep)

        detections = tf.concat([final_roi, final_ids[:, None], final_probs[:, None]],
                               axis=1)
        padding = tf.maximum(self.num_detect-tf.shape(detections)[0], 0)
        detections = tf.pad(detections, [(0, padding), (0, 0)])
        return detections







