import cv2
import time
from tqdm import tqdm
from ultralytics import YOLO
from ultralytics.yolo.utils.plotting import Annotator
from d_face import face_detection
# from dec_duration import detect_duration

# model_police = YOLO('models/police0508.pt')
# model_electromobile = YOLO('models/electromobile0509.pt')
# model_coco = YOLO(r"E:\pythonProject\xznsh\scrfd-opencv\weights\best20230606.pt")
# action_model = YOLO(r'E:\pythonProject\xznsh\scrfd-opencv\weights\action_recognition.pt')

def analysis_video(source_path,output_path,people_modle_path,face_modle_path,action_modle_path):
# def analysis_video(source_path, output_path):
    model_coco = YOLO(people_modle_path)
    action_model = YOLO(action_modle_path)
    # start_time = time.time()
    cap = cv2.VideoCapture(source_path)

    # 直接从视频的第 frameToStart 帧开始
    frameToStart = 1430
    cap.set(cv2.CAP_PROP_POS_FRAMES, frameToStart)

    # total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    # tqdm_desc = 'Processing video frames'
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    output_movie = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    if frameToStart != 0:
        count = frameToStart
    else:
        count = 0
    # 标记有没有出现过拉门的动作
    action_flag = {
        "action":0,
        "action_frame":[]
        }
    # 标记有没有出现过人脸
    face_flag = {
        "face":0,
        "frame":[]
        }
    # 定义帧数字典
    XJ_dict = {
        "head":0,
        "tail":0
        }
    # with tqdm(total=total_frames, desc=tqdm_desc) as pbar:
    while cap.isOpened():
        # Read a frame from the video
        success, frame = cap.read()
        # cv2.imshow("f",frame)
        # cv2.waitKey(0)

        count += 1
        if success:
            # 第一步：用COCO数据集推理
            results_coco = model_coco(frame)
            action_result = action_model(frame)
            # print("*"*100, results_coco.masks)
            
            for r in results_coco:
                annotator = Annotator(frame, line_width=1)   
                # print("*"*100, r.boxes)
                boxes = r.boxes
                for box in boxes:                             
                    b = box.xyxy[0]  # get box coordinates in (x1,y1,x2,y2) format  #tensor([ 677.5757,  147.2737, 1182.3381,  707.2565])
                    b_i = b.int() + 1
                    c = box.cls  # tensor([0.])
                    confidence = float(box.conf)
                    confidence = round(confidence, 2)
                    # 过滤置信度0.5以下目标
                    if confidence < 0.5:
                        continue
                    if c.int() == 1:
                        if XJ_dict['head'] == 0 :
                            XJ_dict['head'] = count
                        else:
                            XJ_dict['tail'] = count
                        # print ("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!巡检人员")
                        crop_img = frame[b_i[1]:b_i[3],b_i[0]:b_i[2]]
                        frame = face_detection(face_modle_path,frame,crop_img,b_i[0],b_i[1],b_i[2],b_i[3],face_flag,count)
                        # annotator = Annotator(frame, line_width=1)
                    annotator.box_label(b, model_coco.names[int(c)]+str(confidence),(0,0,255))

            for r_a in action_result:
                annotator_a = Annotator(frame, line_width=1)    
                # print("*"*100, r.boxes)
                boxes_a = r_a.boxes
                if boxes_a:
                    action_flag["action"] += 1
                    action_flag["action_frame"].append(count)
                for box_a in boxes_a:                             
                    b_a = box_a.xyxy[0]  # get box coordinates in (x1,y1,x2,y2) format  #tensor([ 677.5757,  147.2737, 1182.3381,  707.2565])
                    # b_i_a = b_a.int() + 1
                    c_a = box_a.cls  # tensor([0.])
                    confidence_a = float(box_a.conf)
                    confidence_a = round(confidence_a, 2)
                    # 过滤置信度0.5以下目标
                    if confidence_a < 0.5:
                        continue
                    # annotator = Annotator(frame, line_width=1)
                    annotator_a.box_label(b_a, action_model.names[int(c_a)]+str(confidence_a),(255,0,0))
                    
            annotated_frame_coco = annotator.result()
            annotated_a_frame_coco = annotator_a.result()
            output_movie.write(annotated_a_frame_coco)
            # pbar.update(1)
        else:
            # Break the loop if the end of the video is reached
            break
    
    cap.release()
    output_movie.release()

    diff = round((XJ_dict["tail"]-XJ_dict["head"])/fps,2)
    fina_frame = [round(_ /fps,2) for _ in face_flag["frame"]] 
    s = ', '.join(map(str, fina_frame))
    action_frame = [round(_ /fps,2) for _ in action_flag["action_frame"]] 
    s_action = ', '.join(map(str, action_frame))

    return diff,face_flag,s,action_flag,s_action

# if __name__ == '__main__':
#     analysis_video(r"E:\pythonProject\xznsh\scrfd-opencv\data\xj_video.mp4", r"E:\pythonProject\xznsh\scrfd-opencv\result\video.mp4")