import cv2
from tqdm import tqdm
from ultralytics import YOLO
from ultralytics.yolo.utils.plotting import Annotator
import os

import cv2

import time
import os
import queue
import threading

from pipeline_atm_getResult import atm_det
from pipeline_atm_anaysis import analysis_pph,analysis_button


class analysis_atm():

    def __init__(self,imgPath,savePath,modellist):

        self.imgPath = imgPath
        self.savePath = savePath

        self.imgList = os.listdir(self.imgPath)

        #加载模型
        self.model_person = YOLO(modellist[0])
        self.model_pp_hand = YOLO(modellist[1])
        self.model_blue = YOLO(modellist[2])
        self.model_screen = YOLO(modellist[3])

        # 图片检测后队列
        self.imgQueue1 = queue.Queue(maxsize=0)
        self.imgQueue2 = queue.Queue(maxsize=0)

        #线程
        self.get_img_listThread = threading.Thread(target=self.get_img_list)
        self.analysis_personThread = threading.Thread(target=self.analysis_person)
        self.draw_imagesThread = threading.Thread(target=self.draw_images)


    # 读图
    def get_img_list(self):

        for image_path in self.imgList:

            imagepath = os.path.join(self.imgPath,image_path) 

            self.imgQueue1.put(imagepath)       

    # 检测图
    def analysis_person(self):

        count = 0

        while True:
            if ~self.imgQueue1.empty():

                imagepath = self.imgQueue1.get()

                count = count + 1

                print('-----------------------------------------start',count,self.imgPath,'---------------------------------------------------------')
                # 检测工作人员信息
                score_person = atm_det.get_person_result(imagepath,self.model_person)
                # 检测当前屏幕信息
                score_screen = atm_det.get_screen_result(imagepath,self.model_screen)

                print('-----------------------------------------score_person:',score_person,'score_screen:',score_screen,'---------------------------------------------------------')
                

                # 屏幕信息列表
                screen_pph_list = ['Sign','Enter_Password_1','Enter_Password_2']
                screen_password_list = ['Enter_Password_1','Enter_Password_2']
                # screen_button_list = ['Select_Service_Item','Bank_Card','ID_Card',"Face_Verification"]

                imglist = []

                if score_person and score_screen:

                    #检测屏幕中按钮情况
                    score_button = atm_det.get_blue_result(imagepath,self.model_blue)

                    # print('score_button:',score_button)
                    
                    # 屏幕编号
                    screen_id  = list(score_screen[0].keys())[0]
                    # 检测人员类别
                    person_labels = list(score_person[0].keys())[0]
                    person_ppox = list(score_person[0].values())[0]

                    # print('screen_id:',screen_id,'person_labels:',person_labels)

                    #工作人员行为检测列表
                    # person_re_list = []

                    # 屏幕+工作人员
                    screen_dict = {screen_id  + "_wrong":list(score_screen[0].values())[0]}

                    # person_re_list.append(screen_dict)

                    person_dict = {person_labels + "_wrong":person_ppox}
                    
                    # person_re_list.append(person_dict)
                    person_re_list = [screen_dict,person_dict]

                    # print('person_re_list:',person_re_list)


                        
                    # # 检测到工作人员状态再检测签字情况
                    # if person_labels == 'wrong' and screen_id == 'Sign' :
                    if person_labels == 'wrong' and screen_id == 'Sign':
                    # if person_labels == 'wrong' :

                        score_pp_hand = analysis_pph(imagepath,self.model_pp_hand,score_person)

                        if score_pp_hand:

                            pp_hand_dict = {list(score_pp_hand[0].keys())[0] + '_hand_dict_wrong' :list(score_pp_hand[0].values())[0]}
                            person_re_list.append(pp_hand_dict)
                        else:
                            pp_hand_dict = {'_hand_dict_right':[50,50]}
                            # pp_hand_dict = {'pp_hand_dict_right':'sign and enter password is right action.'}
                            person_re_list.append(pp_hand_dict)

                    # 检测输入密码环节
                    if person_labels == 'wrong' and screen_id in screen_password_list:
                    # if person_labels == 'wrong' :

                        score_pp_hand = analysis_pph(imagepath,self.model_pp_hand,score_person)

                        if list(score_pp_hand[0].keys())[0] == 'password_area':

                            pp_hand_dict = {list(score_pp_hand[0].keys())[0] + '_pp_dict_wrong' :list(score_pp_hand[0].values())[0]}
                            person_re_list.append(pp_hand_dict)
                        else:
                            pp_hand_dict = {'pp_dict_right':[50,50]}
                            # pp_hand_dict = {'pp_hand_dict_right':'sign and enter password is right action.'}
                            person_re_list.append(pp_hand_dict)


                    # 检测屏幕中按钮与首部情况情况
                    if score_button:

                        # 读取图片，判断按钮的位置

                        # print('score_button:',score_button)

                        button_result = analysis_button(imagepath,score_button,score_person)

                        if button_result:
                            button_dict ={'button_dict_wrong':button_result[0]}
                            person_re_list.append(button_dict)
                        else:
                            # 判别顾客的手误识别
                            button_dict ={'button_dict_right':[60,60]}
                            # button_dict ={'button_dict_right':'button action is right.'}
                            person_re_list.append(button_dict)

                    


                    # 检测人员回避情况
                    if screen_id in screen_password_list:

                        if person_labels == 'wrong' or 'correct':

                            avoid_dict = {'avoid_dict_wrong':person_ppox}
                            person_re_list.append(avoid_dict)
                        else:
                            # print('person_labels:',person_labels)
                            avoid_dict = {'avoid_dict_right':[70,70]}
                            # avoid_dict = {'avoid_dict_right':person_ppox}
                            person_re_list.append(avoid_dict) 

                    # print('person_re_list',person_re_list)

                    print('person_re_list_2:',person_re_list)

                # else:
                #     pass

                # print('person_re_list_2:',person_re_list)
                
                # imglist.append(person_re_list)
                
                    img_dict = {imagepath:person_re_list}
                    print('img_dict:',img_dict)

                    self.imgQueue2.put(img_dict)




    def draw_images(self):   

        while True:
            if ~self.imgQueue2.empty():

                imagedict = self.imgQueue2.get()

                print('---------------------------------------','imagedict:',imagedict,'---------------------------------------')
                imgpath = list(imagedict.keys())[0]
                img_result = list(imagedict.values())[0]

                imgcv = cv2.imread(imgpath)
                imgname = os.path.basename(imgpath)

                print('img_result:',img_result)


                for re_dic in img_result:

                    # print(re_dic)

                    re_txt = list(re_dic.keys())[0]

                    re_right = re_txt.split('_')[-1]
                    re_bbox = list(re_dic.values())[0]

                    print('re_bbox:',re_bbox)

                    print('---------------------------------------',re_right,'---------------------------------------')

                    if re_right == "wrong":

                        # re_bbox = list(re_dic.values())[0]

                        cv2.rectangle(imgcv, (int(re_bbox[0]), int(re_bbox[1])),(int(re_bbox[2]), int(re_bbox[3])), (0, 0, 255), 1)

                        cv2.putText(imgcv, re_txt, (int(re_bbox[0]) - 10, int(re_bbox[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 1)

                        # cv2.imwrite(os.path.join(self.savePath,imgname),imgcv)
                    else:

                        print('re_txt:',re_txt)

                        cv2.putText(imgcv, re_txt, (int(re_bbox[0]), int(re_bbox[1])), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255, 0, 255), 3)

                cv2.imwrite(os.path.join(self.savePath,imgname),imgcv)



    def run(self):

        self.get_img_listThread.start()
        self.analysis_personThread.start()
        self.draw_imagesThread.start()         





if __name__ == '__main__':
    



    modelList = ("model_files/bk1.pt","model_files/best_pph.pt","model_files/best_butten_h_01.pt","model_files/best_screen_02.pt")

    # imgPath= 'E:/BANK_XZ_yolov8/data_file'
    imgPath = 'E:/Images Data/XZ/DatasetId_1845281_1685322163/Images'
    # imgPath = 'E:/Images Data/XZ/DatasetId_1845279_1685326217/DatasetId_1845279_1685326217/Images'
    savepath = 'E:/BANK_XZ_yolov8/output_data_06'

    # imgList = os.listdir(imgPath)

    # for image_path in imgList:

    #     imagepath = os.path.join(imgPath,image_path)

    #     # imagepath = 'E:/BANK_XZ_yolov8/data_file/0001326.jpg'

    a = analysis_atm(imgPath,savepath,modelList)
    a.run()