# BSD 3-Clause License
#
# Copyright (c) 2017 xxxx
# All rights reserved.
# Copyright 2021 Huawei Technologies Co., Ltd
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# * Neither the name of the copyright holder nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
# ============================================================================
"""
This script compiles all the bounding boxes in the training data and
clusters them for each convout resolution on which they're used.

Run this script from the Yolact root directory.
"""

import os.path as osp
import json, pickle
import sys
from math import sqrt
from itertools import product
import torch
import random

import numpy as np

dump_file = 'weights/bboxes.pkl'
aug_file  = 'weights/bboxes_aug.pkl'

use_augmented_boxes = True


def intersect(box_a, box_b):
    """ We resize both tensors to [A,B,2] without new malloc:
    [A,2] -> [A,1,2] -> [A,B,2]
    [B,2] -> [1,B,2] -> [A,B,2]
    Then we compute the area of intersect between box_a and box_b.
    Args:
      box_a: (tensor) bounding boxes, Shape: [A,4].
      box_b: (tensor) bounding boxes, Shape: [B,4].
    Return:
      (tensor) intersection area, Shape: [A,B].
    """
    A = box_a.size(0)
    B = box_b.size(0)
    max_xy = torch.min(box_a[:, 2:].unsqueeze(1).expand(A, B, 2),
                       box_b[:, 2:].unsqueeze(0).expand(A, B, 2))
    min_xy = torch.max(box_a[:, :2].unsqueeze(1).expand(A, B, 2),
                       box_b[:, :2].unsqueeze(0).expand(A, B, 2))
    inter = torch.clamp((max_xy - min_xy), min=0)
    return inter[:, :, 0] * inter[:, :, 1]


def jaccard(box_a, box_b, iscrowd=False):
    """Compute the jaccard overlap of two sets of boxes.  The jaccard overlap
    is simply the intersection over union of two boxes.  Here we operate on
    ground truth boxes and default boxes. If iscrowd=True, put the crowd in box_b.
    E.g.:
        A ∩ B / A ∪ B = A ∩ B / (area(A) + area(B) - A ∩ B)
    Args:
        box_a: (tensor) Ground truth bounding boxes, Shape: [num_objects,4]
        box_b: (tensor) Prior boxes from priorbox layers, Shape: [num_priors,4]
    Return:
        jaccard overlap: (tensor) Shape: [box_a.size(0), box_b.size(0)]
    """
    inter = intersect(box_a, box_b)
    area_a = ((box_a[:, 2]-box_a[:, 0]) *
              (box_a[:, 3]-box_a[:, 1])).unsqueeze(1).expand_as(inter)  # [A,B]
    area_b = ((box_b[:, 2]-box_b[:, 0]) *
              (box_b[:, 3]-box_b[:, 1])).unsqueeze(0).expand_as(inter)  # [A,B]
    union = area_a + area_b - inter

    if iscrowd:
        return inter / area_a
    else:
        return inter / union  # [A,B]

# Also convert to point form
def to_relative(bboxes):
    return np.concatenate((bboxes[:, 2:4] / bboxes[:, :2], (bboxes[:, 2:4] + bboxes[:, 4:]) / bboxes[:, :2]), axis=1)


def make_priors(conv_size, scales, aspect_ratios):
    prior_data = []
    conv_h = conv_size[0]
    conv_w = conv_size[1]

    # Iteration order is important (it has to sync up with the convout)
    for j, i in product(range(conv_h), range(conv_w)):
        x = (i + 0.5) / conv_w
        y = (j + 0.5) / conv_h
        
        for scale, ars in zip(scales, aspect_ratios):
            for ar in ars:
                w = scale * ar / conv_w
                h = scale / ar / conv_h

                # Point form
                prior_data += [x - w/2, y - h/2, x + w/2, y + h/2]
    
    return np.array(prior_data).reshape(-1, 4)

# fixed_ssd_config
# scales = [[3.5, 4.95], [3.6, 4.90], [3.3, 4.02], [2.7, 3.10], [2.1, 2.37], [2.1, 2.37], [1.8, 1.92]]
# aspect_ratios = [ [[1, sqrt(2), 1/sqrt(2), sqrt(3), 1/sqrt(3)][:n], [1]] for n in [3, 5, 5, 5, 3, 3, 3] ]
# conv_sizes = [(35, 35), (18, 18), (9, 9), (5, 5), (3, 3), (2, 2)]

scales = [[1.68, 2.91],
          [2.95, 2.22, 0.84],
          [2.23, 2.17, 3.12],
          [0.76, 1.94, 2.72],
          [2.10, 2.65],
          [1.80, 1.92]]
aspect_ratios = [[[0.72, 0.96], [0.68, 1.17]],
                 [[1.28, 0.66], [0.63, 1.23], [0.89, 1.40]],
                 [[2.05, 1.24], [0.57, 0.83], [0.61, 1.15]],
                 [[1.00, 2.21], [0.47, 1.60], [1.44, 0.79]],
                 [[1.00, 1.41, 0.71, 1.73, 0.58], [1.08]],
                 [[1.00, 1.41, 0.71, 1.73, 0.58], [1.00]]]
conv_sizes = [(35, 35), (18, 18), (9, 9), (5, 5), (3, 3), (2, 2)]

# yrm33_config
# scales = [ [5.3] ] * 5
# aspect_ratios = [ [[1, 1/sqrt(2), sqrt(2)]] ]*5
# conv_sizes = [(136, 136), (67, 67), (33, 33), (16, 16), (8, 8)]


SMALL = 0
MEDIUM = 1
LARGE = 2

if __name__ == '__main__':
        
    with open(dump_file, 'rb') as f:
        bboxes = pickle.load(f)

    sizes = []
    smalls = []
    for i in range(len(bboxes)):
        area = bboxes[i][4] * bboxes[i][5]
        if area < 32 ** 2:
            sizes.append(SMALL)
            smalls.append(area)
        elif area < 96 ** 2:
            sizes.append(MEDIUM)
        else:
            sizes.append(LARGE)

    # Each box is in the form [im_w, im_h, pos_x, pos_y, size_x, size_y]
    
    if use_augmented_boxes:
        with open(aug_file, 'rb') as f:
            bboxes_rel = pickle.load(f)
    else:
        bboxes_rel = to_relative(np.array(bboxes))
        

    with torch.no_grad():
        sizes = torch.Tensor(sizes)

        anchors = [make_priors(cs, s, ar) for cs, s, ar in zip(conv_sizes, scales, aspect_ratios)]
        anchors = np.concatenate(anchors, axis=0)
        anchors = torch.Tensor(anchors).cuda()

        bboxes_rel = torch.Tensor(bboxes_rel).cuda()
        perGTAnchorMax = torch.zeros(bboxes_rel.shape[0]).cuda()

        chunk_size = 1000
        for i in range((bboxes_rel.size(0) // chunk_size) + 1):
            start = i * chunk_size
            end   = min((i + 1) * chunk_size, bboxes_rel.size(0))
            
            ious = jaccard(bboxes_rel[start:end, :], anchors)
            maxes, maxidx = torch.max(ious, dim=1)

            perGTAnchorMax[start:end] = maxes
    

        hits = (perGTAnchorMax > 0.5).float()

        print('Total recall: %.2f' % (torch.sum(hits) / hits.size(0) * 100))
        print()

        for i, metric in zip(range(3), ('small', 'medium', 'large')):
            _hits = hits[sizes == i]
            _size = (1 if _hits.size(0) == 0 else _hits.size(0))
            print(metric + ' recall: %.2f' % ((torch.sum(_hits) / _size) * 100))