import numpy as np

# 对于第一张feature_map, 产生anchor box的代码如下：
feature_size = (38, 38)
offsets = (0.5, 0.5)
step = 8
size = (30, 60)
ratios = [1, 2, 0.5]

anchors = []
for i in range(feature_size[0]):
    for j in range(feature_size[1]):
        cy = (i + offsets[0]) * step
        cx = (j + offsets[1]) * step
        min = size[0]
        max = np.sqrt(size[0] * size[1])
        anchors.append([cx, cy, min, max])
        anchors.append([cx, cy, max, max])
        for r in ratios[1:]:
            sr = np.sqrt(r)
            w = min * sr
            h = min / sr
            anchors.append([cx, cy, w, h])
# anchors = np.array(anchors).reshape(-1, 4)
anchors = np.array(anchors).reshape(1, 1, feature_size[0], feature_size[1], -1)

print(anchors.shape)
print(anchors[:8])

# Feature1: 38*38,每个像素点4个anchor,总共5776个anchor(步长为8)
# 每个像素点4个anchor, (0,0)位置对应anchor如下:
# w/h=1:   [4, 4, 30, 30]
#          [4, 4, 42.42640687,  42.42640687]
# w/h=2:   [4, 4, 42.42640687,  21.21320344]
# w/h=0.5: [4, 4, 21.21320344,  42.42640687]

class SSDTargetGenerator(mx.gluon.Block):
    def __init__(self, iou_thresh=0.5, neg_thresh=0.5, negative_mining_ratio=3,
                 stds=[0.1, 0.1, 0.2, 0.2], **kwargs):
        super(SSDTargetGenerator, self).__init__()
        self.iou_thresh = iou_thresh
        self.stds = mx.nd.array(stds)


    def forward(self, anchors, cls_preds, gt_boxes, gt_ids):
        ‘‘‘
        Parameters
        ----------
        anchors: num_anchor*4
        cls_preds:None
        gt_boxes: 1*num_gt*4
        gt_ids: 1*num_gt*1

        Returns
        -------
        cls_targets: 1*num_anchor
        box_targets: 1*num_anchor*4

        ‘‘‘
        # print(1, anchors.shape, gt_boxes.shape, gt_ids.shape)
        anchors_cornors = self._center_to_corner(anchors)
        ious = self.box_iou(gt_boxes, anchors_cornors)  #[num_gt, num_anchor]
        # [num_gt,]
        best_prior_overlap, best_prior_idx = ious.max(axis=1), ious.argmax(axis=1)
        #[num_anchor,]
        best_truth_overlap, best_truth_idx = ious.max(axis=0), ious.argmax(axis=0)  #每一个anchor与它IOU最大的gt_box匹配
        best_truth_overlap[best_prior_idx] = 2  # ensure best prior    #与gt_box的IOU最大的anchor，将其IOU设置为2，保证每个gt_box与其IOU最大的anchor匹配
        # ensure every gt matches with its prior of max overlap
        for j in range(best_prior_idx.shape[0]):   #与gt_box的IOU最大的anchor，将best_truth_idx中该anchor的匹配索引设置为对应gt_box
            best_truth_idx[best_prior_idx[j]] = j

        matches = gt_boxes[0, best_truth_idx]
        cls_targets = gt_ids[0, best_truth_idx] + 1   #num_anchor*1
        cls_targets[best_truth_overlap < self.iou_thresh] = 0   #num_anchor*4, 若anchor与gt_box的IOU小于阈值，设置为不匹配
        box_targets = self.encode(matches, anchors, self.stds)
        # print(3, cls_targets.shape, box_targets.shape)
        return cls_targets.reshape(1,-1), box_targets.reshape(1, -1, 4)

    def _center_to_corner(self, anchors):
        boxes = mx.nd.concat(anchors[:, :2] - anchors[:, 2:]/2, anchors[:, :2] + anchors[:, 2:]/2, dim=1)
        return boxes

    def box_iou(self, box_a, box_b):
        assert box_a.shape[0] == 1
        box_a = box_a[0]
        A = box_a.shape[0]
        B = box_b.shape[0]

        # s1 = box_b[:, 2:].repeat(repeats=A, axis=0).reshape(B, A, 2).transpose(axes=(1, 0, 2))
        # print(s1.shape)
        max_xy = mx.nd.minimum(box_a[:, 2:].repeat(repeats=B, axis=0).reshape(A, B, 2),
                           box_b[:, 2:].repeat(repeats=A, axis=0).reshape(B, A, 2).transpose(axes=(1, 0, 2)))
        min_xy = mx.nd.maximum(box_a[:, :2].repeat(repeats=B, axis=0).reshape(A, B, 2),
                           box_b[:, :2].repeat(repeats=A, axis=0).reshape(B, A, 2).transpose(axes=(1, 0, 2)))
        diff = (max_xy-min_xy)
        inter = mx.nd.clip(diff, a_min=0, a_max=mx.nd.max(diff).asscalar())
        inter_area = inter[:, :, 0] * inter[:, :, 1]
        area_a = (box_a[:, 2] - box_a[:, 0])*(box_a[:, 3] - box_a[:, 1])
        area_a = area_a.repeat(repeats=B, axis=0).reshape(A, B)
        area_b = (box_b[:, 2] - box_b[:, 0])*(box_b[:, 3] - box_b[:, 1])
        area_b = area_b.repeat(repeats=A, axis=0).reshape(B, A).transpose(axes=(1, 0))
        # print(inter_area.shape, area_a.shape, area_b.shape)
        return inter_area/(area_a+area_b-inter_area)

    def encode(self, matches, priors, stds):
        # encode variance
        # print(2, matches.shape, priors.shape)
        g_cxcy = (matches[:, :2] + matches[:, 2:])/2 - priors[:, :2]
        g_cxcy = g_cxcy/stds[:2]

        g_wh = (matches[:, 2:]-matches[:, :2])/priors[:, 2:]
        g_wh = mx.nd.log(g_wh)/stds[2:]
        return mx.nd.concat(g_cxcy, g_wh, dim=1)

class SSDMultiBoxLoss(gluon.Block):
    r"""Single-Shot Multibox Object Detection Loss.

    .. note::

        Since cross device synchronization is required to compute batch-wise statistics,
        it is slightly sub-optimal compared with non-sync version. However, we find this
        is better for converged model performance.

    Parameters
    ----------
    negative_mining_ratio : float, default is 3
        Ratio of negative vs. positive samples.
    rho : float, default is 1.0
        Threshold for trimmed mean estimator. This is the smooth parameter for the
        L1-L2 transition.
    lambd : float, default is 1.0
        Relative weight between classification and box regression loss.
        The overall loss is computed as :math:`L = loss_class + lambda 	imes loss_loc`.
    min_hard_negatives : int, default is 0
        Minimum number of negatives samples.

    """
    def __init__(self, negative_mining_ratio=3, rho=1.0, lambd=1.0,
                 min_hard_negatives=0, **kwargs):
        super(SSDMultiBoxLoss, self).__init__(**kwargs)
        self._negative_mining_ratio = max(0, negative_mining_ratio)
        self._rho = rho
        self._lambd = lambd
        self._min_hard_negatives = max(0, min_hard_negatives)

    def forward(self, cls_pred, box_pred, cls_target, box_target):
        """Compute loss in entire batch across devices."""
        # require results across different devices at this time
        cls_pred, box_pred, cls_target, box_target = [_as_list(x)             for x in (cls_pred, box_pred, cls_target, box_target)]
        # cross device reduction to obtain positive samples in entire batch
        num_pos = []
        for cp, bp, ct, bt in zip(*[cls_pred, box_pred, cls_target, box_target]):
            pos_samples = (ct > 0)
            num_pos.append(pos_samples.sum())
        num_pos_all = sum([p.asscalar() for p in num_pos])
        if num_pos_all < 1 and self._min_hard_negatives < 1:
            # no positive samples and no hard negatives, return dummy losses
            cls_losses = [nd.sum(cp * 0) for cp in cls_pred]
            box_losses = [nd.sum(bp * 0) for bp in box_pred]
            sum_losses = [nd.sum(cp * 0) + nd.sum(bp * 0) for cp, bp in zip(cls_pred, box_pred)]
            return sum_losses, cls_losses, box_losses


        # compute element-wise cross entropy loss and sort, then perform negative mining
        cls_losses = []
        box_losses = []
        sum_losses = []
        for cp, bp, ct, bt in zip(*[cls_pred, box_pred, cls_target, box_target]):
            pred = nd.log_softmax(cp, axis=-1)
            pos = ct > 0
            cls_loss = -nd.pick(pred, ct, axis=-1, keepdims=False)
            rank = (cls_loss * (pos - 1)).argsort(axis=1).argsort(axis=1)     # loss进行排序
            hard_negative = rank < nd.maximum(self._min_hard_negatives, pos.sum(axis=1) 
                                              * self._negative_mining_ratio).expand_dims(-1)  #选择多少个负样本，作为比较难的负样本
            # mask out if not positive or negative
            cls_loss = nd.where((pos + hard_negative) > 0, cls_loss, nd.zeros_like(cls_loss))   #正样本+选择的负样本，参与loss计算
            cls_losses.append(nd.sum(cls_loss, axis=0, exclude=True) / max(1., num_pos_all))

            bp = _reshape_like(nd, bp, bt)
            box_loss = nd.abs(bp - bt)
            box_loss = nd.where(box_loss > self._rho, box_loss - 0.5 * self._rho,
                                (0.5 / self._rho) * nd.square(box_loss))
            # box loss only apply to positive samples
            box_loss = box_loss * pos.expand_dims(axis=-1)   #只有正样本有box_loss, 负样本的box_loss置0
            box_losses.append(nd.sum(box_loss, axis=0, exclude=True) / max(1., num_pos_all))
            sum_losses.append(cls_losses[-1] + self._lambd * box_losses[-1])

        return sum_losses, cls_losses, box_losses