import matplotlib.pyplot as plt
import cv2
import numpy as np

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

def draw_det(img, results):
    ## draw bbox
    for obj in results:

        x = int(obj[0])
        y = int(obj[1])
        x2 = int(obj[2])
        y2 = int(obj[3])
        
        color = (0,0,0)

        tk = 12
        text_size = 4
        
        img = cv2.rectangle(img, (x,y), (x2,y2), color, tk) # bbox

        # content =  obj['cls_name']
        # img = cv2.putText(img, content, (x, y-10), font, text_size, color, text_size, cv2.LINE_AA) # label
    
    return img

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

        x = int(obj['bbox'][0])
        y = int(obj['bbox'][1])
        x2 = int(obj['bbox'][2])
        y2 = int(obj['bbox'][3])
        
        color = (_COLORS[obj['cls_id']%8] * 255).astype(np.uint8).tolist()

        tk = 12
        text_size = 4
        
        img = cv2.rectangle(img, (x,y), (x2,y2), color, tk) # bbox

        content =  obj['cls_name']
        img = cv2.putText(img, content, (x, y-10), font, text_size, color, text_size, cv2.LINE_AA) # label
    
    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):
    ## draw bbox
    for obj in worker_results:

        x = int(obj['bbox'][0])
        y = int(obj['bbox'][1])
        x2 = int(obj['bbox'][2])
        y2 = int(obj['bbox'][3])
        
        color = (_COLORS[obj['cls_id']%2+10] * 255).astype(np.uint8).tolist()

        tk = 12
        text_size = 4
        
        img = cv2.rectangle(img, (x,y), (x2,y2), color, tk) # bbox

        content =  obj['cls_name']
        img = cv2.putText(img, content, (x, y-10), font, text_size, color, text_size, 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


_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)