from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import cv2


vcoco_action_string = {2: 'hold', 3: 'stand', 4: 'sit', 5: 'ride', 6: 'walk',\
                       7: 'look', 8: 'hit_inst', 9: 'hit_obj', 10: 'eat_obj', \
                       11: 'eat_inst', 12: 'jump', 13: 'lay', 14: 'talk', 15: \
                       'carry', 16: 'throw', 17: 'catch', 18: 'cut_inst', 19:'cut_obj', \
                       20: 'run', 21: 'work_on_comp', 22: 'ski', 23: 'surf', 24: 'skateboard', \
                       25: 'smile', 26: 'drink', 27: 'kick', 28: 'point', 29: 'read', 30: 'snowboard'}
def draw_box_on_img(box, img,color=None):

    vis_img = img.copy()
    box = [int(x) for x in box]
    cv2.rectangle(vis_img, (box[0], box[1]), (box[2], box[3]), color, 2)
    draw_point=[int((box[0]+box[2])*1.0/2),int((box[1]+box[3])*1.0/2)]
    
    return vis_img,color


def draw_line_on_img_vcoco(box,line, img, class_index,color):
    
    vis_img = img.copy()
    font=cv2.FONT_HERSHEY_SIMPLEX
    x=int(box[0])+2
    y=int(box[1])+2
    f=int(box[1])+2
    for i in range(len(class_index)):
     
        font_scale=1
        font_thickness=2
    
        text_size, _ = cv2.getTextSize(vcoco_action_string[class_index[i]] , font, font_scale, font_thickness)
        vis_img=cv2.rectangle(vis_img,(x,y),(x+text_size[0],y+text_size[1]+5),color[1],-1)
        

        vis_img=cv2.putText(vis_img, vcoco_action_string[class_index[i]] ,(x,y + text_size[1] ),font,font_scale,[51,255,153],font_thickness)
        y=y+text_size[1]+5
   
    return vis_img,y


def draw_img_vcoco(img, output_i, top_k,threshold,color):
    list_action = []
    for action in output_i['hoi_prediction']:
        subject_id = action['subject_id']
        object_id = action['object_id']
        category_id = action['category_id']
        score = action['score']
        single_out = [subject_id,object_id,category_id,score]
        list_action.append(single_out)
    list_action = sorted(list_action, key=lambda x:x[-1], reverse=True)
    action_dict = []
    action_cate = []
    action_color=[]
    subj_box=[]
    sb={}
    sbj=[]
    for action in list_action[:top_k]:

        subject_id,object_id,category_id,score = action
        if score<threshold:
            break
        subject_obj = output_i['predictions'][subject_id]
        subject_box = subject_obj['bbox']
        object_obj = output_i['predictions'][object_id]
        object_box = object_obj['bbox']

        point_1 = [int((subject_box[0]+subject_box[2])*1.0/2),int((subject_box[1]+subject_box[3])*1.0/2)]
        point_2 = [int((object_box[0]+object_box[2])*1.0/2),int((object_box[1]+object_box[3])*1.0/2)]
  
        if [point_1,point_2] not in action_dict:
            
            img,color_hum = draw_box_on_img(subject_box, img, color[subject_obj['category_id']]['color'])
            
            img,color_obj = draw_box_on_img(object_box, img, color[object_obj['category_id']]['color'])
            
            action_dict.append([point_1,point_2])
            action_color.append([color_hum,color_obj])
            subj_box.append([int(subject_box[0]),int(subject_box[1])])    
           
        action_cate.append([])
        action_cate[action_dict.index([point_1,point_2])].append(category_id)

    for i,(action_item,clr) in enumerate(zip(action_dict,action_color)):
        
        img,offset = draw_line_on_img_vcoco(subj_box[i],action_item,img,action_cate[action_dict.index(action_item)],clr)

        for p in range(i+1,len(subj_box)):
            if subj_box[p]==subj_box[i]:
                subj_box[p][1]=offset
    return img