# Copyright 2021 The KaiJIN Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import numpy as np
from itertools import product as product
import torch
from torch.autograd import Function


def nms_(dets, thresh):
  """
  Courtesy of Ross Girshick
  [https://github.com/rbgirshick/py-faster-rcnn/blob/master/lib/nms/py_cpu_nms.py]
  """
  x1 = dets[:, 0]
  y1 = dets[:, 1]
  x2 = dets[:, 2]
  y2 = dets[:, 3]
  scores = dets[:, 4]

  areas = (x2 - x1) * (y2 - y1)
  order = scores.argsort()[::-1]

  keep = []
  while order.size > 0:
    i = order[0]
    keep.append(int(i))
    xx1 = np.maximum(x1[i], x1[order[1:]])
    yy1 = np.maximum(y1[i], y1[order[1:]])
    xx2 = np.minimum(x2[i], x2[order[1:]])
    yy2 = np.minimum(y2[i], y2[order[1:]])

    w = np.maximum(0.0, xx2 - xx1)
    h = np.maximum(0.0, yy2 - yy1)
    inter = w * h
    ovr = inter / (areas[i] + areas[order[1:]] - inter)

    inds = np.where(ovr <= thresh)[0]
    order = order[inds + 1]

  return np.array(keep).astype(np.int)


def decode(loc, priors, variances):
  boxes = torch.cat((
      priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
      priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)
  boxes[:, :2] -= boxes[:, 2:] / 2
  boxes[:, 2:] += boxes[:, :2]
  return boxes


def nms(boxes, scores, overlap=0.5, top_k=200):
  keep = scores.new(scores.size(0)).zero_().long()
  if boxes.numel() == 0:
    return keep, 0
  x1 = boxes[:, 0]
  y1 = boxes[:, 1]
  x2 = boxes[:, 2]
  y2 = boxes[:, 3]
  area = torch.mul(x2 - x1, y2 - y1)
  v, idx = scores.sort(0)  # sort in ascending order
  # I = I[v >= 0.01]
  idx = idx[-top_k:]  # indices of the top-k largest vals
  xx1 = boxes.new()
  yy1 = boxes.new()
  xx2 = boxes.new()
  yy2 = boxes.new()
  w = boxes.new()
  h = boxes.new()

  # keep = torch.Tensor()
  count = 0
  while idx.numel() > 0:
    i = idx[-1]  # index of current largest val
    # keep.append(i)
    keep[count] = i
    count += 1
    if idx.size(0) == 1:
      break
    idx = idx[:-1]  # remove kept element from view
    # load bboxes of next highest vals
    torch.index_select(x1, 0, idx, out=xx1)
    torch.index_select(y1, 0, idx, out=yy1)
    torch.index_select(x2, 0, idx, out=xx2)
    torch.index_select(y2, 0, idx, out=yy2)
    # store element-wise max with next highest score
    xx1 = torch.clamp(xx1, min=x1[i])
    yy1 = torch.clamp(yy1, min=y1[i])
    xx2 = torch.clamp(xx2, max=x2[i])
    yy2 = torch.clamp(yy2, max=y2[i])
    w.resize_as_(xx2)
    h.resize_as_(yy2)
    w = xx2 - xx1
    h = yy2 - yy1
    # check sizes of xx1 and xx2.. after each iteration
    w = torch.clamp(w, min=0.0)
    h = torch.clamp(h, min=0.0)
    inter = w * h
    # IoU = i / (area(a) + area(b) - i)
    rem_areas = torch.index_select(area, 0, idx)  # load remaining areas)
    union = (rem_areas - inter) + area[i]
    IoU = inter / union  # store result in iou
    # keep only elements with an IoU <= overlap
    idx = idx[IoU.le(overlap)]
  return keep, count


class PriorBox(object):
  """Compute priorbox coordinates in center-offset form for each source
  feature map.
  """

  def __init__(self, input_size, feature_maps):
    super(PriorBox, self).__init__()
    self.imh = input_size[0]
    self.imw = input_size[1]

    # number of priors for feature map location (either 4 or 6)
    self.variance = [0.1, 0.2]
    #self.feature_maps = cfg.FEATURE_MAPS
    self.min_sizes = [16, 32, 64, 128, 256, 512]
    self.steps = [4, 8, 16, 32, 64, 128]
    self.clip = False
    self.feature_maps = feature_maps

  def forward(self):
    mean = []
    for k in range(len(self.feature_maps)):
      feath = self.feature_maps[k][0]
      featw = self.feature_maps[k][1]
      for i, j in product(range(feath), range(featw)):
        f_kw = self.imw / self.steps[k]
        f_kh = self.imh / self.steps[k]

        cx = (j + 0.5) / f_kw
        cy = (i + 0.5) / f_kh

        s_kw = self.min_sizes[k] / self.imw
        s_kh = self.min_sizes[k] / self.imh

        mean += [cx, cy, s_kw, s_kh]

    output = torch.Tensor(mean).view(-1, 4)
    if self.clip:
      output.clamp_(max=1, min=0)
    return output


class Detect(Function):
  """At test time, Detect is the final layer of SSD.  Decode location preds,
  apply non-maximum suppression to location predictions based on conf
  scores and threshold to a top_k number of output predictions for both
  confidence score and locations.
  """

  def __init__(self):
    self.num_classes = 2
    self.top_k = 5000
    self.nms_thresh = 0.3
    self.conf_thresh = 0.05
    self.variance = [0.1, 0.2]

  def forward(self, loc_data, conf_data, prior_data):
    """
    Args:
        loc_data: (tensor) Loc preds from loc layers
            Shape: [batch,num_priors*4]
        conf_data: (tensor) Shape: Conf preds from conf layers
            Shape: [batch*num_priors,num_classes]
        prior_data: (tensor) Prior boxes and variances from priorbox layers
            Shape: [1,num_priors,4] 
    """
    num = loc_data.size(0)
    num_priors = prior_data.size(0)

    conf_preds = conf_data.view(
        num, num_priors, self.num_classes).transpose(2, 1)
    batch_priors = prior_data.view(-1, num_priors,
                                   4).expand(num, num_priors, 4)
    batch_priors = batch_priors.contiguous().view(-1, 4)

    decoded_boxes = decode(loc_data.view(-1, 4),
                           batch_priors, self.variance)
    decoded_boxes = decoded_boxes.view(num, num_priors, 4)

    output = torch.zeros(num, self.num_classes, self.top_k, 5)

    for i in range(num):
      boxes = decoded_boxes[i].clone()
      conf_scores = conf_preds[i].clone()

      for cl in range(1, self.num_classes):
        c_mask = conf_scores[cl].gt(self.conf_thresh)
        scores = conf_scores[cl][c_mask]

        if scores.dim() == 0:
          continue
        l_mask = c_mask.unsqueeze(1).expand_as(boxes)
        boxes_ = boxes[l_mask].view(-1, 4)
        ids, count = nms(
            boxes_, scores, self.nms_thresh, self.top_k)
        output[i, cl, :count] = torch.cat((scores[ids[:count]].unsqueeze(1),
                                           boxes_[ids[:count]]), 1)

    return output
