import asyncio
import glob
from argparse import ArgumentParser
from tqdm import tqdm
import matplotlib.pyplot as plt
import cv2
import numpy as np
import json
import torch

from mmcv import Config
from mmdet.apis import async_inference_detector, inference_detector, show_result_pyplot
from mmdet.apis.inference import init_detector

from ensemble_boxes import *

def get_json_data(label_path):
    with open(label_path, encoding='utf-8-sig') as f:
        json_data = json.load(f)
    return json_data

# Based on BGR

##################### 1. Draw Vehicle  ##########################
font =  cv2.FONT_HERSHEY_PLAIN

def draw_vehicle(img, vehicle_results):
    ## draw bbox
    count = 0
    for obj in vehicle_results:

        x = int(obj['bbox'][0])
        y = int(obj['bbox'][1])
        w = int(obj['bbox'][2])
        h = int(obj['bbox'][3])

        trk_id = int(obj['trk_id'])
        color = (_COLORS[trk_id%70] * 255).astype(np.uint8).tolist()

        tk = 3
        # if obj['risk1'] == 'danger':
        #     tk = 5

        img = cv2.rectangle(img, (x,y), (x+w,y+h), color, tk) # bbox

#         content =  str(obj['trk_id'])+'_'+obj['cls_name']

#         img = cv2.putText(img, content, (x, y-10), font, 2, color, 2, cv2.LINE_AA) # label


    ## draw boundary & risk 
    for obj in vehicle_results:

        x = int(obj['bbox'][0])
        y = int(obj['bbox'][1])
        w = int(obj['bbox'][2])
        h = int(obj['bbox'][3])

        trk_id = int(obj['trk_id'])
        color = (_COLORS[trk_id%70] * 255).astype(np.uint8).tolist()

        tk = 2
        if obj['risk1'] == 'danger':
        #     tk = 5
            # img = cv2.putText(img, "No signalman", (x+10, y+40), font, 3, color, 3, cv2.LINE_AA)
            img = cv2.putText(img, "WARNING! No signalman", (100, 100), font, 4, (0,0,255), 4, cv2.LINE_AA)
            count +=1

        # left
        img = cv2.line(img, (x-int(w/2),y), (x-int(w/2),y+h), color, tk)
        img = cv2.line(img, (x-w,y), (x-w,y+h), color, 1)

        # right
        img = cv2.line(img, (x+int(1.5*w),y), (x+int(1.5*w),y+h), color, tk)     
        img = cv2.line(img, (x+2*w,y), (x+2*w,y+h), color, 1) 
    
    # if count == 0:
    #     img = cv2.putText(img, "Safe", (100, 100), font, 4, (0,0,255), 4, cv2.LINE_AA)
    
    return img

def draw_vehicle2(img, vehicle_results):
    ## draw bbox
    color = (0,212,255)
    for obj in vehicle_results:

        x = int(obj['bbox'][0])
        y = int(obj['bbox'][1])
        w = int(obj['bbox'][2])
        h = int(obj['bbox'][3])
      
        tk = 3
        img = cv2.rectangle(img, (x,y), (x+w,y+h), color, tk) # bbox
        
        content =  str(obj['det_score'])+'_'+obj['cls_name']

        img = cv2.putText(img, content, (x, y-10), font, 2, color, 2, cv2.LINE_AA) # label


    ## draw boundary & risk 
    for obj in vehicle_results:

        x = int(obj['bbox'][0])
        y = int(obj['bbox'][1])
        w = int(obj['bbox'][2])
        h = int(obj['bbox'][3])

        tk = 2
        if obj['risk1'] == 'danger':
            img = cv2.putText(img, "WARNING! No signalman", (100, 100), font, 4, (0,0,255), 4, cv2.LINE_AA)

        # left
        img = cv2.line(img, (x-int(w/2),y), (x-int(w/2),y+h), color, tk)
        img = cv2.line(img, (x-w,y), (x-w,y+h), color, 1)

        # right
        img = cv2.line(img, (x+int(1.5*w),y), (x+int(1.5*w),y+h), color, tk)     
        img = cv2.line(img, (x+2*w,y), (x+2*w,y+h), color, 1) 
    
    
    return img

##################################################################

##################### 2. Draw Worker  ############################
def draw_worker(img, worker_results):
    
    if worker_results is not None:
        ## draw bbox
        for obj in worker_results:
            x = int(obj['bbox'][0])
            y = int(obj['bbox'][1])
            w = int(obj['bbox'][2])
            h = int(obj['bbox'][3])

            tk = 3
            color = (255,255,255)

            if obj['cls_name'] == 'signalman':
                color = (0,255,0)
            else:
                if obj['risk1'] == 'danger':
                    color = (0,0,255)


            img = cv2.rectangle(img, (x,y), (x+w,y+h), color, tk) # bbox

            content =  str(obj['det_score'])

            img = cv2.putText(img, content, (x, y-10), font, tk, color, tk, cv2.LINE_AA) # label
                                               
    return img
                                               
##################################################################


                                               
                            
def risk_algorithm_1(vehicle_results, worker_results):
    for idx,obj in enumerate(vehicle_results): # based on X-axis
        left_bound = obj['bbox'][0]-obj['bbox'][2]
        right_bound = obj['bbox'][0]+obj['bbox'][2]*2
        top_bound = obj['bbox'][1]
        bottom_bound = obj['bbox'][1]+obj['bbox'][3]
        
        count = 0
        for jdx,obj in enumerate(worker_results):
            if obj['cls_name']=='signalman':
                count+=1
            else:
                box = obj['bbox']
                if (box[0]+(box[2]/2))>left_bound and (box[0]+(box[2]/2))<right_bound and (box[1]+(box[3]/2))>top_bound and (box[1]+(box[3]/2))<bottom_bound:
                    worker_results[jdx]['risk1'] = 'danger'

        
        if count>0:
            vehicle_results[idx]['risk1'] = 'safe'
            
        else:
            vehicle_results[idx]['risk1'] = 'danger'
        
    return vehicle_results, worker_results




# for ensemble
def convert_wbf(results,threshold,image_size):
    boxes_list = []
    scores_list = []
    labels_list = []
    for predicted_class in range(len(results)):
        for preds in results[predicted_class]:
            if preds[-1] >= threshold:
                # print(preds)
                box = preds[:4]
                boxes_list.append([box[0]/image_size[1],
                                   box[1]/image_size[0],
                                   box[2]/image_size[1],
                                   box[3]/image_size[0]]
                                   )
                score = preds[-1].astype(float)
                scores_list.append(score)
                labels_list.append(predicted_class+1)
    return boxes_list, labels_list, scores_list

# Segmentation modules
def segment_image(cropped_img):
    img = cv2.bilateralFilter(cropped_img,9,105,105)
    r,g,b=cv2.split(img)
    equalize1= cv2.equalizeHist(r)
    equalize2= cv2.equalizeHist(g)
    equalize3= cv2.equalizeHist(b)
    equalize=cv2.merge((r,g,b))
    equalize = cv2.cvtColor(equalize,cv2.COLOR_RGB2GRAY)
    ret,thresh_image = cv2.threshold(equalize,0,255,cv2.THRESH_OTSU+cv2.THRESH_BINARY)
    equalize= cv2.equalizeHist(thresh_image)
    return equalize

def check_signalman(worker_bboxes, frame):

    temp = []
    for x,row in enumerate(worker_bboxes):
        cropped_img = frame[int(row[1]): int(row[1]+row[3]), 
                            int(row[0]+(row[2]/4)): int(row[0]+(row[2]*3/4))]
        img_hsv=cv2.cvtColor(cropped_img, cv2.COLOR_BGR2HSV)
        img_rgb = cropped_img.copy()
        h,w,c=img_rgb.shape

        # cropped to half
        cropped_size = cropped_img.shape
        half_cropped_img = cropped_img[0: int(cropped_size[0]/2), 
                                    0: int(cropped_size[1])]
        # segment
        img_rgb = half_cropped_img.copy()
        segmented_img_rgb = segment_image(img_rgb)
        updated_equalize = segmented_img_rgb.copy()
        updated_equalize[updated_equalize == 0] = 1
        updated_equalize[updated_equalize == 255] = 0
        updated_equalize_rgb = cv2.cvtColor(updated_equalize,cv2.COLOR_GRAY2RGB)
        matched_img = np.multiply(updated_equalize_rgb, img_rgb)
        img_rgb = matched_img
        h,w,c=img_rgb.shape

        # count red
        count = 0
        for idx,i in enumerate(img_rgb):
            for jdx,j in enumerate(i):
                if j[0]>j[1]+20 and j[0]>j[2]+20:
                    img_rgb[idx][jdx][0]=255
                    img_rgb[idx][jdx][1]=0
                    img_rgb[idx][jdx][2]=0
                    count +=1
        if count/(h*w)>=0.04 and count/(h*w)<0.7:
            temp.append([row[0],row[1],row[2],row[3],2]) # xywh scores cls
        else:
            temp.append([row[0],row[1],row[2],row[3],1])
    return temp

# x y w h -> cx cy w h
def yolo_coor(box):
    center_x = int(box[0] + (box[2]/2))
    center_y = int(box[1] + (box[3]/2))
    w = int(box[2])
    h = int(box[3])
    
    return [center_x,center_y,w,h]

def fp16_clamp(x, min=None, max=None):
    if not x.is_cuda and x.dtype == torch.float16:
        # clamp for cpu float16, tensor fp16 has no clamp implementation
        return x.float().clamp(min, max).half()

    return x.clamp(min, max)

# input x1,y1,x2,y2
def bbox_iou(bboxes1,bboxes2):
    area1 = (bboxes1[..., 2] - bboxes1[..., 0]) * (
        bboxes1[..., 3] - bboxes1[..., 1])
    area2 = (bboxes2[..., 2] - bboxes2[..., 0]) * (
        bboxes2[..., 3] - bboxes2[..., 1])
    
    lt = torch.max(bboxes1[..., :2], bboxes2[..., :2])  # [B, rows, 2]
    rb = torch.min(bboxes1[..., 2:], bboxes2[..., 2:])  # [B, rows, 2]
    
    wh = fp16_clamp(rb - lt, min=0)
    overlap = wh[..., 0] * wh[..., 1]
    
    union = area1 + area2 - overlap
    
    eps=1e-6
    eps = union.new_tensor([eps])
    union = torch.max(union, eps)
    ious = overlap / union
    
    return ious

_COLORS = np.array(
    [   0.000, 0.447, 0.741,
        0.850, 0.325, 0.098,
        0.494, 0.184, 0.556,
        0.466, 0.674, 0.188,
        0.301, 0.745, 0.933,
        0.635, 0.078, 0.184,
        0.749, 0.749, 0.000,
        0.667, 0.000, 1.000,
        0.333, 0.333, 0.000,
        0.333, 0.667, 0.000,
        0.333, 1.000, 0.000,
        0.667, 0.333, 0.000,
        0.667, 0.667, 0.000,
        0.667, 1.000, 0.000,
        1.000, 0.333, 0.000,
        1.000, 0.667, 0.000,
        1.000, 1.000, 0.000,
        0.000, 0.333, 0.500,
        0.000, 0.667, 0.500,
        0.000, 1.000, 0.500,
        0.333, 0.000, 0.500,
        0.333, 0.333, 0.500,
        0.333, 0.667, 0.500,
        0.333, 1.000, 0.500,
        0.667, 0.000, 0.500,
        0.667, 0.333, 0.500,
        0.667, 0.667, 0.500,
        0.667, 1.000, 0.500,
        1.000, 0.000, 0.500,
        1.000, 0.333, 0.500,
        1.000, 0.667, 0.500,
        1.000, 1.000, 0.500,
        0.000, 0.333, 1.000,
        0.000, 0.667, 1.000,
        0.000, 1.000, 1.000,
        0.333, 0.000, 1.000,
        0.333, 0.333, 1.000,
        0.333, 0.667, 1.000,
        0.333, 1.000, 1.000,
        0.667, 0.000, 1.000,
        0.667, 0.333, 1.000,
        0.667, 0.667, 1.000,
        0.667, 1.000, 1.000,
        1.000, 0.000, 1.000,
        1.000, 0.333, 1.000,
        1.000, 0.667, 1.000,
        0.333, 0.000, 0.000,
        0.500, 0.000, 0.000,
        0.667, 0.000, 0.000,
        0.833, 0.000, 0.000,
        1.000, 0.000, 0.000,
        0.000, 0.167, 0.000,
        0.000, 0.333, 0.000,
        0.000, 0.500, 0.000,
        0.000, 0.667, 0.000,
        0.000, 0.833, 0.000,
        0.000, 1.000, 0.000,
        0.000, 0.000, 0.167,
        0.000, 0.000, 0.333,
        0.000, 0.000, 0.500,
        0.000, 0.000, 0.667,
        0.000, 0.000, 0.833,
        0.000, 0.000, 1.000,
        0.000, 0.000, 0.000,
        0.143, 0.143, 0.143,
        0.286, 0.286, 0.286,
        0.429, 0.429, 0.429,
        0.571, 0.571, 0.571,
        0.714, 0.714, 0.714,
        0.857, 0.857, 0.857,
        0.000, 0.447, 0.741,
        0.314, 0.717, 0.741,
        0.50, 0.5, 0]).astype(np.float32).reshape(-1, 3)