import re
import sys
import os
import time
import math
import numpy as np

import itertools
import struct  # get_image_size
import imghdr  # get_image_size

ENDS_WITH_NUMBER = re.compile(r'(.*)_[0-9]+')

def get_darknet_files(weights_file):
  """
  return an appropriate weights, cfg, and names file if they exist.
  Returns:
    (str, str, str, str): weights, cfg, names, prefix for other files
  """
  
  if not weights_file.endswith(".weights"):
    weights_file += ".weights"

  prefix = os.path.splitext(weights_file)[0]
  original = prefix
  
  cfg_file = f"{prefix}.cfg"
  names_file = f"{prefix}.names"

  if os.path.isfile(weights_file):
    if not os.path.isfile(cfg_file) or not os.path.isfile(names_file):
      
      if prefix.endswith("_best"):
        prefix = prefix.replace("_best", "")

      elif prefix.endswith("_last"):
        prefix = prefix.replace("_last", "")

      elif prefix.endswith("_final"):
        prefix = prefix.replace("_final", "")

      else:
        m = ENDS_WITH_NUMBER.match(prefix)
        
        if m:
          prefix = m.group(1)
      
      if original != prefix:
        cfg_file = f"{prefix}.cfg"
        names_file = f"{prefix}.names"
      
  if not (os.path.isfile(cfg_file) and os.path.isfile(cfg_file) and
    os.path.isfile(weights_file)):
  
    cfg_file = None
    names_file = None
    weights_file = None
  
  return weights_file, cfg_file, names_file, prefix

def sigmoid(x):
  return 1.0 / (np.exp(-x) + 1.)


def softmax(x):
  x = np.exp(x - np.expand_dims(np.max(x, axis=1), axis=1))
  x = x / np.expand_dims(x.sum(axis=1), axis=1)
  return x


def bbox_iou(box1, box2, x1y1x2y2=True):
  
  # print('iou box1:', box1)
  # print('iou box2:', box2)

  if x1y1x2y2:
    mx = min(box1[0], box2[0])
    Mx = max(box1[2], box2[2])
    my = min(box1[1], box2[1])
    My = max(box1[3], box2[3])
    w1 = box1[2] - box1[0]
    h1 = box1[3] - box1[1]
    w2 = box2[2] - box2[0]
    h2 = box2[3] - box2[1]
  else:
    w1 = box1[2]
    h1 = box1[3]
    w2 = box2[2]
    h2 = box2[3]

    mx = min(box1[0], box2[0])
    Mx = max(box1[0] + w1, box2[0] + w2)
    my = min(box1[1], box2[1])
    My = max(box1[1] + h1, box2[1] + h2)
  uw = Mx - mx
  uh = My - my
  cw = w1 + w2 - uw
  ch = h1 + h2 - uh
  carea = 0
  if cw <= 0 or ch <= 0:
    return 0.0

  area1 = w1 * h1
  area2 = w2 * h2
  carea = cw * ch
  uarea = area1 + area2 - carea
  return carea / uarea


def nms_cpu(boxes, confs, nms_thresh=0.5, min_mode=False):
  # print(boxes.shape)
  x1 = boxes[:, 0]
  y1 = boxes[:, 1]
  x2 = boxes[:, 2]
  y2 = boxes[:, 3]

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

  keep = []
  while order.size > 0:
    idx_self = order[0]
    idx_other = order[1:]

    keep.append(idx_self)

    xx1 = np.maximum(x1[idx_self], x1[idx_other])
    yy1 = np.maximum(y1[idx_self], y1[idx_other])
    xx2 = np.minimum(x2[idx_self], x2[idx_other])
    yy2 = np.minimum(y2[idx_self], y2[idx_other])

    w = np.maximum(0.0, xx2 - xx1)
    h = np.maximum(0.0, yy2 - yy1)
    inter = w * h

    if min_mode:
      over = inter / np.minimum(areas[order[0]], areas[order[1:]])
    else:
      over = inter / (areas[order[0]] + areas[order[1:]] - inter)

    inds = np.where(over <= nms_thresh)[0]
    order = order[inds + 1]
  
  return np.array(keep)



def plot_boxes_cv2(img, boxes, savename=None, class_names=None, color=None):
  import cv2
  img = np.copy(img)
  colors = np.array([[1, 0, 1], [0, 0, 1], [0, 1, 1], [0, 1, 0], [1, 1, 0], [1, 0, 0]], dtype=np.float32)

  def get_color(c, x, max_val):
    ratio = float(x) / max_val * 5
    i = int(math.floor(ratio))
    j = int(math.ceil(ratio))
    ratio = ratio - i
    r = (1 - ratio) * colors[i][c] + ratio * colors[j][c]
    return int(r * 255)

  width = img.shape[1]
  height = img.shape[0]
  for i in range(len(boxes)):
    box = boxes[i]
    x1 = int(box[0] * width)
    y1 = int(box[1] * height)
    x2 = int(box[2] * width)
    y2 = int(box[3] * height)
    bbox_thick = int(0.6 * (height + width) / 600)
    if color:
      rgb = color
    else:
      rgb = (255, 0, 0)
    if len(box) >= 7 and class_names:
      cls_conf = box[5]
      cls_id = box[6]
      print(f"{class_names[cls_id]} ({cls_conf}) at coords: {x1},{y1}->{x2},{y2}")
      classes = len(class_names)
      offset = cls_id * 123457 % classes
      red = get_color(2, offset, classes)
      green = get_color(1, offset, classes)
      blue = get_color(0, offset, classes)
      if color is None:
        rgb = (red, green, blue)
      msg = str(class_names[cls_id])+" "+str(round(cls_conf,3))
      t_size = cv2.getTextSize(msg, 0, 0.7, thickness=bbox_thick // 2)[0]
      c1, c2 = (x1,y1), (x2, y2)
      c3 = (c1[0] + t_size[0], c1[1] - t_size[1] - 3)
      print(f"c1={c1}, c2={c2}, c3={c3}, rgb={rgb}, ")
      cv2.rectangle(img, (x1,y1), (c3[0], c3[1]), rgb, -1)
      #cv2.rectangle(img, (x1,y1), (np.float32(c3[0]), np.float32(c3[1])), rgb, -1)
      img = cv2.putText(img, msg, (c1[0], c1[1] - 2), cv2.FONT_HERSHEY_SIMPLEX,0.7, (0,0,0), bbox_thick//2,lineType=cv2.LINE_AA)
    
    img = cv2.rectangle(img, (x1, y1), (x2, y2), rgb, bbox_thick)
  if savename:
    #print("save plot results to %s" % savename)
    cv2.imwrite(savename, img)
  return img


def read_truths(lab_path):
  if not os.path.exists(lab_path):
    return np.array([])
  if os.path.getsize(lab_path):
    truths = np.loadtxt(lab_path)
    truths = truths.reshape(truths.size / 5, 5)  # to avoid single truth problem
    return truths
  else:
    return np.array([])


def load_class_names(namesfile):
  class_names = []
  with open(namesfile, 'r') as fp:
    lines = fp.readlines()
  for line in lines:
    line = line.rstrip()
    class_names.append(line)
  return class_names



def post_processing(img, conf_thresh, nms_thresh, output, debug=False):

  # anchors = [12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401]
  # num_anchors = 9
  # anchor_masks = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
  # strides = [8, 16, 32]
  # anchor_step = len(anchors) // num_anchors

  # [batch, num, 1, 4]
  box_array = output[0]
  # [batch, num, num_classes]
  confs = output[1]

  t1 = time.time()

  if type(box_array).__name__ != 'ndarray':
    box_array = box_array.cpu().detach().numpy()
    confs = confs.cpu().detach().numpy()

  num_classes = confs.shape[2]

  if debug:
    print(f"post_processing: using num_classes={num_classes}")

  # [batch, num, 4]
  box_array = box_array[:, :, 0]

  # [batch, num, num_classes] --> [batch, num]
  max_conf = np.max(confs, axis=2)
  max_id = np.argmax(confs, axis=2)

  t2 = time.time()

  bboxes_batch = []
  for i in range(box_array.shape[0]):
     
    argwhere = max_conf[i] > conf_thresh
    l_box_array = box_array[i, argwhere, :]
    l_max_conf = max_conf[i, argwhere]
    l_max_id = max_id[i, argwhere]

    bboxes = []
    # nms for each class
    for j in range(num_classes):

      cls_argwhere = l_max_id == j
      ll_box_array = l_box_array[cls_argwhere, :]
      ll_max_conf = l_max_conf[cls_argwhere]
      ll_max_id = l_max_id[cls_argwhere]

      keep = nms_cpu(ll_box_array, ll_max_conf, nms_thresh)
      
      if (keep.size > 0):
        ll_box_array = ll_box_array[keep, :]
        ll_max_conf = ll_max_conf[keep]
        ll_max_id = ll_max_id[keep]

        for k in range(ll_box_array.shape[0]):
          bboxes.append([ll_box_array[k, 0], ll_box_array[k, 1], ll_box_array[k, 2], ll_box_array[k, 3], ll_max_conf[k], ll_max_conf[k], ll_max_id[k]])
    
    bboxes_batch.append(bboxes)

  t3 = time.time()

  if debug:
    print('-----------------------------------')
    print('       max and argmax : %f' % (t2 - t1))
    print('                  nms : %f' % (t3 - t2))
    print('Post processing total : %f' % (t3 - t1))
    print('-----------------------------------')
  
  return bboxes_batch
