import os
import numpy as np
import xml.etree.ElementTree as ET
from xml.dom import minidom
from pathlib import Path
import math

import cv2


def load_ground_truth(anno_path):
    ground_truth_dict = {}
    for root, dirs, files in os.walk(anno_path):
        for f in files:
            frame = int(f.split('.xml')[0])
            xml_path = os.path.join(root, f)
            tree = ET.parse(xml_path)
            ground_truth_dict[frame] = []
            for item in tree.findall(".//object"):
                xmax = float(item.find(".//xmax").text)
                xmin = float(item.find(".//xmin").text)
                ymax = float(item.find(".//ymax").text)
                ymin = float(item.find(".//ymin").text)
                track_id = int(item.find(".//trackid").text)
                box_item = {}
                box_item['xmax'] = xmax
                box_item['xmin'] = xmin
                box_item['ymax'] = ymax
                box_item['ymin'] = ymin
                box_item['track_id'] = track_id
                ground_truth_dict[frame].append(box_item)

    return ground_truth_dict


if __name__ == "__main__":
    video_dir = "/Users/wujiahang/Project/EdgeTile/EdgeTileClient/data/drone/video_2k_120fps/"
    video_name = "uav0000138_00000_v.mp4"
    seq_dir = Path('/Users/wujiahang/Project/EdgeTile/EdgeTileClient/data/drone/anno_2k_120fps/') / Path(video_name).stem
    video_path = Path(video_dir) / Path(video_name)
    cap = cv2.VideoCapture(str(video_path))  # 打开指定路径上的视频文件
    video_width = cap.get(3)
    video_height = cap.get(4)
    
    frame_id = 0

    seg_height = []
    seg_height.append(0.0)
    divide_w = 10
    divide_h = 6
    for i in range(1, divide_h+1):
        seg_height.append(seg_height[i-1] + video_height * 1.0 / divide_h)

    seg_width = []
    seg_width.append(0.0)
    for i in range(1, divide_w+1):
        seg_width.append(seg_width[i-1] + video_width * 1.0 / divide_w)
    print(seg_height)
    print(seg_width)
    ground_truth_dict = load_ground_truth(seq_dir)

    while (True):
        ret, frame = cap.read()
        a20 = False
        a48 = False
        a168 = False
        a173 = False


        if ret == True:
            for y in seg_height:
                cv2.line(frame, (0, int(y)),
                         (int(video_width), int(y)), (0, 0, 255), 2)
                # print(y)
            for x in seg_width:
                cv2.line(frame, (int(x), 0),
                         (int(x), int(video_height)), (0, 0, 255), 2)
                # print(x)
            grids = []
            for box in ground_truth_dict[frame_id]:
                min_w = 0
                max_w = 0
                min_h = 0
                max_h = 0
                xmin = box['xmin']
                ymin = box['ymin']
                xmax = box['xmax']
                ymax = box['ymax']

                for i in range(divide_w, -1, -1):
                    if(xmin >= seg_width[i]):
                        min_w = i
                        break
                for i in range(0, divide_w+1):
                    if(xmax <= seg_width[i]):
                        max_w = i
                        break
                for i in range(divide_h, -1, -1):
                    if(ymin >= seg_height[i]):
                        min_h = i
                        break
                for i in range(0, divide_h+1):
                    if(ymax <= seg_height[i]):
                        max_h = i
                        break
                for i in range(min_w, max_w):
                    for j in range(min_h, max_h):
                        grids.append(j*divide_w+i)
                        
                # print("out box:[{},{},{},{}]".format(min_w, min_h, max_w,max_h))
                font = cv2.FONT_HERSHEY_SIMPLEX
                imgzi = cv2.putText(frame, '{}'.format(box['track_id']), (int(xmin), int(ymin)), font, 0.8, (0, 0, 255), 2)
                if box['track_id'] == 20:
                    a20 = True
                if box['track_id'] == 20:
                    a48 = True
                if box['track_id'] == 20:
                    a168 = True
                if box['track_id'] == 20:
                    a173 = True
                cv2.rectangle(frame, (int(xmin), int(ymin)),
                              (int(xmax), int(ymax)), (255, 0, 0), 2)
            # if a20 and a48 and a168 and a173:
            #     print('f: {}'.format(frame_id))
            #     input()
            # print(grids)
            grids = list(set(grids))
            # print(grids)
            # print(len(seg_width))
            # print(len(seg_height))
            no_target_grids = []
            for x in range(0, divide_w*divide_h):
                if x not in grids:
                    no_target_grids.append(x)
            # print(no_target_grids)
            for grid in no_target_grids:
                left_top_x = int(grid % divide_w)
                left_top_y = int(grid / divide_w)
                # print("({},{})".format(left_top_x, left_top_y))    
                # sub_img = frame[int(seg_height[left_top_y]):int(seg_height[left_top_y+1]), int(seg_width[left_top_x]):int(seg_width[left_top_x+1])]
                # white_rect = np.ones(sub_img.shape, dtype=np.uint8) * 255
                # res = cv2.addWeighted(frame, 0.5, white_rect, 0.5, 1.0)
                blk = np.zeros(frame.shape, np.uint8)
                cv2.rectangle(blk, (int(seg_width[left_top_x]), int(seg_height[left_top_y])),
                              (int(seg_width[left_top_x+1]), int(seg_height[left_top_y+1])), (255, 255, 255), -1)
                frame = cv2.addWeighted(frame, 1.0, blk, 0.25, 1)
            font = cv2.FONT_HERSHEY_SIMPLEX
            imgzi = cv2.putText(frame, 'frame id:{}'.format(
                frame_id), (50, 300), font, 1.2, (255, 255, 0), 2)
            if(frame_id==550):
                cv2.imwrite("my.jpg",frame)
                cv2.waitKey(0)
            print(frame_id)
            cv2.imshow("frame", frame)
            cv2.waitKey(33)

            frame_id += 1
        else:
            break

    cap.release()
    cv2.destroyAllWindows()