# 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.
# ==============================================================================
"""BoundingBox Coder and Decoder

      Anchors that are generated by AnchorGenerator do not directly be used in
  the detection pipeline. It should two step:

    1. training: encoding bbox groundtruth into generated anchors as network
      prediction targets.
    2. inference: decoding network prediction into anchors.

"""
import abc
import math
import torch
import numpy as np
import tw

#!<-----------------------------------------------------------------------------
#!< Base Box Coder
#!<-----------------------------------------------------------------------------


class _BaseBoxCoder(abc.ABC):
  """BoxCoder is a technical to regress the deltas from network outputs.
  """
  @abc.abstractmethod
  def encode(self, references, proposals):
    """Encode a set of proposals with respect to some reference boxes"""
    raise NotImplementedError

  @abc.abstractmethod
  def decode(self, encodes, anchors):
    """From a set of original boxes and encoded relative box offsets, get the
    decoded boxes."""
    raise NotImplementedError


#!<-----------------------------------------------------------------------------
#!< General Box Coder
#!<-----------------------------------------------------------------------------


class GeneralBoxCoder(_BaseBoxCoder):
  """BoxCoder is a technical to regress the deltas from network outputs.
  """

  def __init__(self,
               means=[0, 0, 0, 0],
               variances=[1, 1, 1, 1],
               bbox_xform_clip=math.log(1000./16),
               clamp_to_image_shape=None):
    super(GeneralBoxCoder, self).__init__()
    self.means = means
    self.variances = variances
    self.bbox_xform_clip = bbox_xform_clip
    self.clamp_to_image_shape = clamp_to_image_shape
    
  def encode(self, references, proposals):
    """references from ground truth, proposals from positive sample.

      dx = (proposal_centered_x - gt_centered_x) / proposal_width
      dy = (proposal_centered_y - gt_centered_y) / proposal_height
      dw = log(gt_width / proposal_width)
      dh = log(gt_height / proposal_height)
      deltas = ([dx, dy, dw, dh] - means) / variances

    Args:
        references ([N, 4(x1, y1, x2, y2)]): groundtruth of bounding box
        proposals ([N, 4(x1, y1, x2, y2)]): anchors of bounding box

    Returns:
        [type]: [deltas (N, 4(&x, &y, &w, &h))]: normalized centered-xywh
    """
    assert proposals.size() == references.size()
    proposals = proposals.float()
    references = references.float()

    # from xyxy to centered-xywh
    px = (proposals[..., 0] + proposals[..., 2]) * 0.5
    py = (proposals[..., 1] + proposals[..., 3]) * 0.5
    pw = proposals[..., 2] - proposals[..., 0] + 1.0
    ph = proposals[..., 3] - proposals[..., 1] + 1.0

    # from xyxy to centered-xywh
    gx = (references[..., 0] + references[..., 2]) * 0.5
    gy = (references[..., 1] + references[..., 3]) * 0.5
    gw = references[..., 2] - references[..., 0] + 1.0
    gh = references[..., 3] - references[..., 1] + 1.0

    # encode
    dx = (gx - px) / pw  # delta(x) / gt-w
    dy = (gy - py) / ph  # delta(y) / gt-h
    dw = torch.log(gw / pw)  # log(w/gt-w)
    dh = torch.log(gh / ph)  # log(h/gt-h)
    deltas = torch.stack([dx, dy, dw, dh], dim=-1)

    means = deltas.new_tensor(self.means).unsqueeze(0)
    variances = deltas.new_tensor(self.variances).unsqueeze(0)
    deltas = deltas.sub_(means).div_(variances)

    return deltas

  def decode(self, encodes, anchors):
    """decode network outputs in terms of anchors. The processing is an inverse
      process compared with encoding.

      denorm_deltas = deltas * variances + means
      denorm_deltas [dx, dy, dw, dh]

      px = (anchor_x1 + anchor_x2) / 2
      py = (anchor_y1 + anchor_y2) / 2
      pw = (anchor_x2 - anchor_x1) + 1
      ph = (anchor_y2 - anchor_y1) + 1

      gx = px + pw * dx
      gy = py + ph * dy
      gw = pw * e^(dw)
      gh = ph * e^(dh)

      x1 = gx - gw * 0.5 + 0.5
      y1 = gy - gh * 0.5 + 0.5
      x2 = gx + gw * 0.5 - 0.5
      y2 = gy + gh * 0.5 - 0.5

      prediction bbox [x1, y1, x2, y2]

    Args:
        encodes ([type]): [network output bbox, (N, 4)]
        anchors ([type]): [anchors, (N, 4)]

    Returns:
        [type]: [prediction bbox (N, 4(x1, y1, x2, y2))]
    """

    means = encodes.new_tensor(self.means).repeat(1, encodes.size(1) // 4)
    variances = encodes.new_tensor(self.variances).repeat(1, encodes.size(1) // 4)
    denorm_deltas = encodes * variances + means

    dx = denorm_deltas[:, 0::4]
    dy = denorm_deltas[:, 1::4]
    dw = denorm_deltas[:, 2::4]
    dh = denorm_deltas[:, 3::4]

    max_ratio = np.abs(np.log(self.bbox_xform_clip))
    dw = dw.clamp(min=-max_ratio, max=max_ratio)
    dh = dh.clamp(min=-max_ratio, max=max_ratio)

    px = ((anchors[:, 0] + anchors[:, 2]) * 0.5).unsqueeze(1).expand_as(dx)
    py = ((anchors[:, 1] + anchors[:, 3]) * 0.5).unsqueeze(1).expand_as(dy)
    pw = (anchors[:, 2] - anchors[:, 0] + 1.0).unsqueeze(1).expand_as(dw)
    ph = (anchors[:, 3] - anchors[:, 1] + 1.0).unsqueeze(1).expand_as(dh)
    gw = pw * dw.exp()
    gh = ph * dh.exp()
    gx = px + pw * dx  # torch.addcmul(px, 1, pw, dx)  # gx = px + pw * dx
    gy = py + ph * dy  # torch.addcmul(py, 1, ph, dy)  # gy = py + ph * dy

    x1 = gx - gw * 0.5 + 0.5
    y1 = gy - gh * 0.5 + 0.5
    x2 = gx + gw * 0.5 - 0.5
    y2 = gy + gh * 0.5 - 0.5

    if self.clamp_to_image_shape is not None:
      x1 = x1.clamp(min=0, max=self.clamp_to_image_shape[1] - 1)
      y1 = y1.clamp(min=0, max=self.clamp_to_image_shape[0] - 1)
      x2 = x2.clamp(min=0, max=self.clamp_to_image_shape[1] - 1)
      y2 = y2.clamp(min=0, max=self.clamp_to_image_shape[0] - 1)

    bboxes = torch.stack([x1, y1, x2, y2], dim=-1).view_as(encodes)
    return bboxes
