import cv2
import os

class Process_tools():

    # 图像文件夹
    def get_video_list(path):
        video_ext = [".mp4", ".avi",".MP4"]
        video_names = []
        for maindir, subdir, file_name_list in os.walk(path):
            for filename in file_name_list:
                apath = os.path.join(maindir, filename)
                ext = os.path.splitext(apath)[1]
                if ext in video_ext:
                    video_names.append(apath)
        return video_names


    # 截取裁剪需要的视频帧
    def save_seg_video(video_name,frameToStart,frametoStop,videoWriter,bbox):

        cap = cv2.VideoCapture(video_name)
        count = 0
        while True:

            success, frame = cap.read()

            if success:

                count += 1
                if count <= frametoStop and count > frameToStart:  # 选取起始帧
                    print('correct= ', count)
                    
                    #裁剪视频画面
                    frame_target = frame[int(bbox[1]):int(bbox[3]), int(bbox[0]):int(bbox[2])]  # (split_height, split_width)

                    videoWriter.write(frame_target)
                    
            if not success or count >= frametoStop:
                break

            print('end')  


    # 获得字典中所有values值（这个值是列表）
    def get_dict_values(lst):
        """
        获取列表中所有字典的 values 值（如果值是列表）
        
        参数:
        lst: 包含字典的列表
        
        返回值:
        values: 包含所有字典的 values 值的列表（如果值是列表）
        """
        return [value for dictionary in lst for value in dictionary.values() if isinstance(value, list)]



    # 解析检测后的结果，为检测后的结果排序
    def analysis_sort_list(result_dict):

        # print('result_dict:',result_dict)

        # 获得检测列表
        re_list = result_dict['start_bbox']
        # print('re_list:',re_list)

        # 获得列表中所有字典的values值
        re_bbox_list = Process_tools.get_dict_values(re_list)

        # 为检测出来的标注框排序
        sorted_lst = sorted(re_bbox_list, key=lambda x: x[0])

        return sorted_lst


    #对比重叠率高的两个部分，并结合标注框，保存最大的标注框
    def contrast_bbox(e_bbox,r_bbox):

        e_bbox_min = e_bbox[:2]
        r_bbox_min = r_bbox[:2]

        bbox_min = [min(x, y) for x, y in zip(e_bbox_min, r_bbox_min)]

        e_bbox_max = e_bbox[-2:]
        r_bbox_max = r_bbox[-2:]

        bbox_max = [max(x, y) for x, y in zip(e_bbox_max, r_bbox_max)]

        bbox = bbox_min + bbox_max

        return bbox



    # 解析result_list列表
    def analysis_re01_list(example_dict,result_dict):

        # 第一次检测到目标的帧率和信息
        example_dict_fps = list(example_dict.keys())[0]
        example_sorted_lst = Process_tools.analysis_sort_list(example_dict)

        # 当前帧检测结果中所有的检测结果数值
        re_dict_fps = list(result_dict.keys())[0]
        re_dict_sorted_lst = Process_tools.analysis_sort_list(result_dict)

        # 保存前后帧率连续的范围、筛选出相同的部分
        cut_list = []
        example_temp = []
        re_temp = []

        for i,ex_bbox in enumerate(example_sorted_lst):

            for j,re_bbox in enumerate(re_dict_sorted_lst):

                iou = Process_tools.calculate_iou(box1=ex_bbox, box2=re_bbox)

                # print(iou)
                
                if iou > 0:

                    bbox = Process_tools.contrast_bbox(e_bbox=ex_bbox,r_bbox=re_bbox)

                    cut_list.append({i:bbox})
                    example_temp.append(ex_bbox)
                    re_temp.append(re_bbox)

                    break

                else:
                    continue

        example_sorted_lst = [item for item in example_sorted_lst if item not in example_temp]
        re_dict_sorted_lst = [item for item in re_dict_sorted_lst if item not in re_temp]

        return cut_list,example_sorted_lst,re_dict_sorted_lst


    # 计算前后帧率重叠范围
    def calculate_iou(box1, box2):
        """
        计算两个边界框之间的IoU值
        
        参数:
        box1: 边界框1的坐标（x1, y1, x2, y2）
        box2: 边界框2的坐标（x1, y1, x2, y2）
        
        返回值:
        iou: 两个边界框之间的IoU值
        """
        x1 = max(box1[0], box2[0])
        y1 = max(box1[1], box2[1])
        x2 = min(box1[2], box2[2])
        y2 = min(box1[3], box2[3])
        
        # 计算交集区域面积
        intersection_area = max(0, x2 - x1 + 1) * max(0, y2 - y1 + 1)
        
        # 计算边界框1和边界框2的面积
        box1_area = (box1[2] - box1[0] + 1) * (box1[3] - box1[1] + 1)
        box2_area = (box2[2] - box2[0] + 1) * (box2[3] - box2[1] + 1)
        
        # 计算并集区域面积
        union_area = box1_area + box2_area - intersection_area
        
        # 计算IoU值
        iou = intersection_area / union_area
        
        return iou

    def para_correction(images_size,bbox,dertpara):

        '''
        修正检测后标注框过小的情况,如果有修正参数则使用修正参数，如果没有就按照坐标值扩大两倍

        '''
        
        if dertpara:
            pass
        else:
            w = (bbox[2] - bbox[0]) /2
            h = (bbox[3] - bbox[1]) /2

            bbox_extand_list_x = [bbox[0] - w,bbox[2] + w]
            bbox_extand_list_y = [bbox[1] - h,bbox[3] + h]
            
            bbox_list_x = Process_tools.contrast(size=images_size[0],bbox_extand_list=bbox_extand_list_x)
            bbox_list_y = Process_tools.contrast(size=images_size[1],bbox_extand_list=bbox_extand_list_y)

            bbox_list = bbox_list_x + bbox_list_y

            return bbox_list
        

    def contrast(size,bbox_extand_list):

        '''
        对比数值是否在这个范围内
        '''
        
        bbox_list = []

        for x in bbox_extand_list:
            
            if x in range(size):
                bbox_list.append(x)
            if x > size:
                bbox_list.append(size)
            if x < 0:
                bbox_list.append(0)
        return bbox_list