import os
import math
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw
import cv2
from tqdm import tqdm

import torch
import torchvision
import torchvision.transforms as transforms
from torchvision.transforms import ToTensor
import timm

from mmdet.apis import async_inference_detector, inference_detector
from mmdet.apis.inference import init_detector

from mmseg.apis import init_model as seg_init_model
from mmseg.apis import inference_model as seg_inference_model
from mmseg.apis import show_result_pyplot as seg_show_result_pyplot
from segment_anything import sam_model_registry, SamPredictor

from mmpose.apis import (inference_topdown, init_model)

# dpt
from torchvision.transforms import Compose
# from dpt.models import DPTDepthModel
# from dpt.midas_net import MidasNet_large
# from dpt.transforms import Resize, NormalizeImage, PrepareForNet

def load_combined_models():
    # Pose estimation model
    pose_config_file = '/media/daitran/Data_SSD1/1_Postdoc_SKKU_2023/Research/manuscript_1_2023/construction-video-to-safety/models/mmpose/configs/body_2d_keypoint/topdown_heatmap/coco/td-hm_hrnet-w32_8xb64-210e_coco-256x192.py'
    pose_checkpoint_file = '/media/daitran/Data_SSD1/1_Postdoc_SKKU_2023/Research/manuscript_1_2023/data/pretrained_models/mmpose/public_topdown/hrnet_w32_coco_256x192-c78dce93_20200708.pth'
    pose_model = init_model(pose_config_file, pose_checkpoint_file, device='cuda:0')
    # Segmentation model
    # segmentation_config_file = '/media/daitran/Data_SSD1/1_Postdoc_SKKU_2023/Research/manuscript_1_2023/construction-video-to-safety/models/mmsegmentation/configs/mask2former/mask2former_swin-l-in22k-384x384-pre_8xb2-160k_ade20k-640x640.py'
    # segmentation_checkpoint = '/media/daitran/Data_SSD1/1_Postdoc_SKKU_2023/Research/manuscript_1_2023/data/pretrained_models/mmseg/public_mask2former/mask2former_swin-l-in22k-384x384-pre_8xb2-160k_ade20k-640x640_20221203_235933-7120c214.pth'
    # semantic_segmantation_model = seg_init_model(segmentation_config_file, segmentation_checkpoint, device='cuda:0')
    # Person detection model
    person_config_file = "/media/daitran/Data_SSD1/1_Postdoc_SKKU_2023/Research/manuscript_1_2023/construction-video-to-safety/models/mmyolo/configs/yolov8/yolov8_x_mask-refine_syncbn_fast_8xb16-500e_coco.py"
    person_checkpoint = '/media/daitran/Data_SSD1/1_Postdoc_SKKU_2023/Research/manuscript_1_2023/data/pretrained_models/mmdet/public_yolov8x/yolov8_x_mask-refine_syncbn_fast_8xb16-500e_coco_20230217_120411-079ca8d1.pth'
    person_model = init_detector(person_config_file, person_checkpoint, device='cuda:0')
    # # Construction vehicle detection model
    # vehicle_config_file = "/media/daitran/Data_SSD1/1_Postdoc_SKKU_2023/Research/manuscript_1_2023/data/pretrained_models/vehicle_detection/yolov8x/yolov8x_vehicle.py"
    # vehicle_checkpoint = '/media/daitran/Data_SSD1/1_Postdoc_SKKU_2023/Research/manuscript_1_2023/data/pretrained_models/vehicle_detection/yolov8x/epoch_49.pth'
    # vehicle_model = init_detector(vehicle_config_file, vehicle_checkpoint, device='cuda:0')
    # # Signalman detection model
    # signalman_config_file = '/media/daitran/Data_SSD1/1_Postdoc_SKKU_2023/Research/manuscript_1_2023/data/pretrained_models/signalman_detection/yolov8x/yolov8-x.py'
    # signalman_checkpoint = '/media/daitran/Data_SSD1/1_Postdoc_SKKU_2023/Research/manuscript_1_2023/data/pretrained_models/signalman_detection/yolov8x/epoch_460.pth'
    # signalman_model = init_detector(signalman_config_file, signalman_checkpoint, device='cuda:0')
    # PPE detection model
    # ppe_config_file = '/media/daitran/Data_SSD1/1_Postdoc_SKKU_2023/Research/manuscript_1_2023/data/pretrained_models/ppe_detection/yolov8x/yolov8x_ppe.py'
    # ppe_checkpoint = '/media/daitran/Data_SSD1/1_Postdoc_SKKU_2023/Research/manuscript_1_2023/data/pretrained_models/ppe_detection/yolov8x/best_coco_bbox_mAP_epoch_11.pth'
    # ppe_model = init_detector(ppe_config_file, ppe_checkpoint, device='cuda:0')
    # # Deptp map estimation model
    # depth_model = torch.hub.load("intel-isl/MiDaS", "MiDaS_small")  # DPT_Hybrid #MiDaS_small
    # depth_model.eval()
    return pose_model, person_model



def mmdet3x_convert_to_bboxes_mmdet(results, threshold, prefix = 'class'):
    boxes_list = []
    scores_list = []
    labels_list = []
    confidence_score = results.pred_instances.scores.tolist()
    for i, conf in enumerate(confidence_score):
        if conf >= threshold:
            # print(conf)
            extracted_box = results.pred_instances.bboxes[i].cpu().tolist()
            extracted_label = results.pred_instances.labels[i].cpu().tolist()
            boxes_list.append([int(extracted_box[0]),
                               int(extracted_box[1]),
                                int(extracted_box[2]),
                                int(extracted_box[3])])
            scores_list.append(conf)
            labels_list.append('{}_{}'.format(prefix,extracted_label+1))
    return boxes_list, labels_list, scores_list
def depth_preprocess_image(image):
    # Resize the input image to dimensions divisible by 32
    h, w, _ = image.shape
    new_h = h - (h % 32)
    new_w = w - (w % 32)
    resized_image = cv2.resize(image, (new_w, new_h))

    input_transform = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    )
    return input_transform(resized_image).unsqueeze(0)
def postprocess_depth(depth_tensor, original_size):
    depth_map = depth_tensor.squeeze().cpu().numpy()
    depth_map = cv2.resize(depth_map, original_size)
    depth_map = depth_map / depth_map.max()
    return depth_map
def extract_depth_map(model, image):
    img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    input_tensor = depth_preprocess_image(img_rgb)
    with torch.no_grad():
        depth_tensor = model(input_tensor)
        depth_map = postprocess_depth(depth_tensor, (img_rgb.shape[1], img_rgb.shape[0]))
    return depth_map
def extract_ankles_points(keypoints, keypoints_confidence):
    keypoint_dict = {}
    keypoint_dict['left_ankle'] = [keypoints[15][0], keypoints[15][1], keypoints_confidence[15]]
    keypoint_dict['right_ankle'] = [keypoints[16][0], keypoints[16][1], keypoints_confidence[16]]
    return keypoint_dict
def extract_keypoints(keypoints):
    left_shoulder= [keypoints[5][0], keypoints[5][1]]
    right_shoulder= [keypoints[6][0], keypoints[6][1]]
    left_hip= [keypoints[11][0], keypoints[11][1]]
    right_hip= [keypoints[12][0], keypoints[12][1]]
    return left_shoulder, right_shoulder, left_hip, right_hip
def estimate_ankle_midpoint(keypoints):
    left_shoulder_coord, right_shoulder_coord, left_hip_coord, right_hip_coord = extract_keypoints(keypoints)
    # Calculate the midpoint of the shoulders and hips
    shoulder_midpoint =[int((left_shoulder_coord[0] + right_shoulder_coord[0]) / 2), int((left_shoulder_coord[1] + right_shoulder_coord[1]) / 2)]
    hip_midpoint = [int((left_hip_coord[0] + right_hip_coord[0]) / 2), int((left_hip_coord[1] + right_hip_coord[1]) / 2)]
    # Calculate the vector from the shoulder midpoint to the hip midpoint
    shoulder_to_hip_vector = [hip_midpoint[0] - shoulder_midpoint[0], hip_midpoint[1] - shoulder_midpoint[1]]
    # # Double the length of the vector to approximate the upper body's mirror image
    mirroring_fator= 1.2
    mirrored_vector = [shoulder_to_hip_vector[0] * mirroring_fator, shoulder_to_hip_vector[1] * mirroring_fator]
    # Add the resulting vector to the hip midpoint to obtain the estimated ankle midpoint
    estimated_ankle_midpoint = [hip_midpoint[0] + mirrored_vector[0], hip_midpoint[1] + mirrored_vector[1]]
    # estimated_ankle_midpoint should be perpendicular with hip_midpoint
    estimated_ankle_midpoint =[hip_midpoint[0],estimated_ankle_midpoint[1]]
    return estimated_ankle_midpoint

def is_point_inside_polygon(x, y, danger_points):
    n = len(danger_points)
    inside = False
    p1x, p1y = danger_points[0]
    for i in range(n + 1):
        p2x, p2y = danger_points[i % n]
        if y > min(p1y, p2y):
            if y <= max(p1y, p2y):
                if x <= max(p1x, p2x):
                    if p1y != p2y:
                        xinters = (y - p1y) * (p2x - p1x) / (p2y - p1y) + p1x
                    if p1x == p2x or x <= xinters:
                        inside = not inside
        p1x, p1y = p2x, p2y
    return inside
def check_inside_polygon(worker_centers, danger_points):
    inside = []
    for i, worker_center in enumerate(worker_centers):
        x, y = worker_center
        inside.append(is_point_inside_polygon(x, y, danger_points))
    return inside
def draw_polygon(frame, danger_points):
    img = Image.fromarray(frame)
    draw = ImageDraw.Draw(img)
    draw.polygon(danger_points, outline ="red", width = 5)
    return np.array(img)

# DPT
def preprocess_img_dpt(image_path):
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) / 255.0
    return img

def depth_estimation_dp(img):
    model_path = "/media/daitran/Data_SSD1/1_Postdoc_SKKU_2023/Research/manuscript_1_2023/data/DPT/weights/dpt_hybrid_nyu-2ce69ec7.pt"
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    net_w = 640
    net_h = 480

    model = DPTDepthModel(
        path=model_path,
        scale=0.000305,
        shift=0.1378,
        invert=True,
        backbone="vitb_rn50_384",
        non_negative=True,
        enable_attention_hooks=False,
    )

    normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
    transform = Compose(
            [
                Resize(
                    net_w,
                    net_h,
                    resize_target=None,
                    keep_aspect_ratio=True,
                    ensure_multiple_of=32,
                    resize_method="minimal",
                    image_interpolation_method=cv2.INTER_CUBIC,
                ),
                normalization,
                PrepareForNet(),
            ]
        )

    model.eval()
    model = model.to(memory_format=torch.channels_last)
    model = model.half()
    model.to(device)

    img_input = transform({"image": img})["image"]
    # computedpt_frame =  cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) / 255.0
    with torch.no_grad():
        sample = torch.from_numpy(img_input).to(device).unsqueeze(0)

        sample = sample.to(memory_format=torch.channels_last)
        sample = sample.half()

        prediction = model.forward(sample)
        prediction = (
            torch.nn.functional.interpolate(
                prediction.unsqueeze(1),
                size=img.shape[:2],
                mode="bicubic",
                align_corners=False,
            )
            .squeeze()
            .cpu()
            .numpy()
        )
    prediction *= 1000.0
    return prediction
