import tensorflow as tf
import numpy as np
import cv2
from configs.kitti_config import config


# ################## From SqueezeDet #########################

def safe_exp(w, thresh):
  """Safe exponential function for tensors."""

  slope = np.exp(thresh)
  with tf.name_scope('safe_exponential'):
    lin_region = tf.to_float(w > thresh)

    lin_out = slope * (w - thresh + 1.)
    exp_out = tf.exp(w)

    out = lin_region * lin_out + (1. - lin_region) * exp_out
  return out


def bbox_transform_inv(bbox):
  """convert a bbox of form [xmin, ymin, xmax, ymax] to [cx, cy, w, h]. Works
  for numpy array or list of tensors.
  """
  with tf.name_scope('bbox_transform_inv') as scope:
    xmin, ymin, xmax, ymax = bbox
    out_box = [[]] * 4

    width = xmax - xmin + 1.0
    height = ymax - ymin + 1.0
    out_box[0] = xmin + 0.5 * width
    out_box[1] = ymin + 0.5 * height
    out_box[2] = width
    out_box[3] = height

  return out_box


def bbox_transform(bbox):
  """convert a bbox of form [cx, cy, w, h] to [xmin, ymin, xmax, ymax]. Works
  for numpy array or list of tensors.
  """
  with tf.name_scope('bbox_transform') as scope:
    cx, cy, w, h = bbox
    out_box = [[]] * 4
    out_box[0] = cx - w / 2
    out_box[1] = cy - h / 2
    out_box[2] = cx + w / 2
    out_box[3] = cy + h / 2

  return out_box


def nms(boxes, probs, threshold):
  """Non-Maximum supression.
  Args:
    boxes: array of [cx, cy, w, h] (center format)
    probs: array of probabilities
    threshold: two boxes are considered overlapping if their IOU is largher than
        this threshold
    form: 'center' or 'diagonal'
  Returns:
    keep: array of True or False.
  """

  order = probs.argsort()[::-1]
  keep = [True]*len(order)

  for i in range(len(order)-1):
    ovps = batch_iou(boxes[order[i+1:]], boxes[order[i]])
    for j, ov in enumerate(ovps):
      if ov > threshold:
        keep[order[j+i+1]] = False
  return keep


def interpre_prediction(prediction, input_mask, anchors, box_input):
  # probability
  batch_size = tf.shape(input_mask)[0]
  num_class_probs = config.NUM_ANCHORS * config.NUM_CLASSES
  pred_class_probs = tf.reshape(
    tf.nn.softmax(
      tf.reshape(
        prediction[:, :, :, :num_class_probs],
        [-1, config.NUM_CLASSES]
      )
    ),
    [batch_size, config.ANCHORS, config.NUM_CLASSES],
    name='pred_class_probs'
  )

  # confidence
  num_confidence_scores = config.NUM_ANCHORS + num_class_probs
  pred_conf = tf.sigmoid(
    tf.reshape(
      prediction[:, :, :, num_class_probs:num_confidence_scores],
      [batch_size, config.ANCHORS]
    ),
    name='pred_confidence_score'
  )

  # bbox_delta
  pred_box_delta = tf.reshape(
    prediction[:, :, :, num_confidence_scores:],
    [batch_size, config.ANCHORS, 4],
    name='bbox_delta'
  )

  # number of object. Used to normalize bbox and classification loss
  num_objects = tf.reduce_sum(input_mask, name='num_objects')

  with tf.name_scope('bbox') as scope:
    with tf.name_scope('stretching'):
      delta_x, delta_y, delta_w, delta_h = tf.unstack(
        pred_box_delta, axis=2)

      anchor_x = anchors[:, 0]
      anchor_y = anchors[:, 1]
      anchor_w = anchors[:, 2]
      anchor_h = anchors[:, 3]

      box_center_x = tf.identity(
        anchor_x + delta_x * anchor_w, name='bbox_cx')
      box_center_y = tf.identity(
        anchor_y + delta_y * anchor_h, name='bbox_cy')
      box_width = tf.identity(
        anchor_w * safe_exp(delta_w, config.EXP_THRESH),
        name='bbox_width')
      box_height = tf.identity(
        anchor_h * safe_exp(delta_h, config.EXP_THRESH),
        name='bbox_height')

    with tf.name_scope('trimming'):
      xmins, ymins, xmaxs, ymaxs = bbox_transform(
        [box_center_x, box_center_y, box_width, box_height])