# 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.
# ============================================================================
"""
Instead of clustering bbox widths and heights, this script
directly optimizes average IoU across the training set given
the specified number of anchor boxes.

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

import pickle
import random
from itertools import product
from math import sqrt

import numpy as np
import torch
from scipy.optimize import minimize

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 torch.Tensor(prior_data).view(-1, 4).cuda()



scales = [[1.68, 2.91], [2.95, 2.22, 0.84], [2.17, 2.22, 3.22], [0.76, 2.06, 2.81], [5.33, 2.79], [13.69]]
aspect_ratios = [[[0.72, 0.96], [0.68, 1.17]], [[1.30, 0.66], [0.63, 1.23], [0.87, 1.41]], [[1.96, 1.23], [0.58, 0.84], [0.61, 1.15]], [[19.79, 2.21], [0.47, 1.76], [1.38, 0.79]], [[4.79, 17.96], [1.04]], [[14.82]]]
conv_sizes = [(35, 35), (18, 18), (9, 9), (5, 5), (3, 3), (2, 2)]

optimize_scales = False

batch_idx = 0


def compute_hits(bboxes, anchors, iou_threshold=0.5):
    ious = jaccard(bboxes, anchors)
    perGTAnchorMax, _ = torch.max(ious, dim=1)
    
    return (perGTAnchorMax > iou_threshold)

def compute_recall(hits, base_hits):
    hits = (hits | base_hits).float()
    return torch.sum(hits) / hits.size(0)


def step(x, x_func, bboxes, base_hits, optim_idx):
    # This should set the scale and aspect ratio
    x_func(x, scales[optim_idx], aspect_ratios[optim_idx])

    anchors = make_priors(conv_sizes[optim_idx], scales[optim_idx], aspect_ratios[optim_idx])

    return -float(compute_recall(compute_hits(bboxes, anchors), base_hits).cpu())


def optimize(full_bboxes, optim_idx, batch_size=5000):
    global batch_idx, scales, aspect_ratios, conv_sizes

    start = batch_idx * batch_size
    end   = min((batch_idx + 1) * batch_size, full_bboxes.size(0))

    if batch_idx > (full_bboxes.size(0) // batch_size):
        batch_idx = 0

    bboxes = full_bboxes[start:end, :]

    anchor_base = [
        make_priors(conv_sizes[idx], scales[idx], aspect_ratios[idx])
            for idx in range(len(conv_sizes)) if idx != optim_idx]
    base_hits = compute_hits(bboxes, torch.cat(anchor_base, dim=0))
    
    
    def set_x(x, scales, aspect_ratios):
        if optimize_scales:
            for i in range(len(scales)):
                scales[i] = max(x[i], 0)
        else:
            k = 0
            for i in range(len(aspect_ratios)):
                for j in range(len(aspect_ratios[i])):
                    aspect_ratios[i][j] = x[k]
                    k += 1
            

    res = minimize(step, x0=scales[optim_idx] if optimize_scales else sum(aspect_ratios[optim_idx], []), method='Powell',
        args = (set_x, bboxes, base_hits, optim_idx),)


def pretty_str(x:list):
    if isinstance(x, list):
        return '[' + ', '.join([pretty_str(y) for y in x]) + ']'
    elif isinstance(x, np.ndarray):
        return pretty_str(list(x))
    else:
        return '%.2f' % x

if __name__ == '__main__':
    
    if use_augmented_boxes:
        with open(aug_file, 'rb') as f:
            bboxes = pickle.load(f)
    else:
        # Load widths and heights from a dump file. Obtain this with
        # python3 scripts/save_bboxes.py
        with open(dump_file, 'rb') as f:
            bboxes = pickle.load(f)

            bboxes = np.array(bboxes)
            bboxes = to_relative(bboxes)

    with torch.no_grad():
        bboxes = torch.Tensor(bboxes).cuda()
        
        def print_out():
            if optimize_scales:
                print('Scales: ' + pretty_str(scales))
            else:
                print('Aspect Ratios: ' + pretty_str(aspect_ratios))

        for p in range(10):
            print('(Sub Iteration) ', end='')
            for i in range(len(conv_sizes)):
                print('%d ' % i, end='', flush=True)
                optimize(bboxes, i)
            print('Done', end='\r')
            
            print('(Iteration %d) ' % p, end='')
            print_out()
            print()

            optimize_scales = not optimize_scales
        
        print('scales = ' + pretty_str(scales))
        print('aspect_ratios = ' + pretty_str(aspect_ratios))