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 mmpose.apis import (inference_topdown, init_model)
from mmpose.apis import (inference_topdown, init_model)
from mmpose.registry import VISUALIZERS
from mmpose.structures import merge_data_samples

from utils import *
import time

video_path = './input/vid_6.mp4'
print("Load models")
pose_model, person_model = load_combined_models()

print("Start processing")
def visualize_all_keypoints(cropped_image):
    cropped_results = inference_topdown(pose_model, cropped_image)
    keypoints = cropped_results[0].pred_instances.keypoints
    keypoints_confidence = cropped_results[0].pred_instances.keypoint_scores
    visualizer = VISUALIZERS.build(pose_model.cfg.visualizer)
    visualizer.set_dataset_meta(pose_model.dataset_meta)
    data_samples = merge_data_samples(cropped_results)
    _ = visualizer.add_datasample(
            'result',
            cropped_image,
            data_sample=data_samples,
            draw_gt=False,
            draw_heatmap=True,
            draw_bbox=False,
            show=False,
            wait_time=0,
            kpt_thr=0.3)
    vis_result = visualizer.get_image()
    return vis_result
    
center_point = (540,360)
radius = 100

trash_threshold = 200
# read video and show frame
cap = cv2.VideoCapture(video_path)
# create video to save processing frame
fps = cap.get(cv2.CAP_PROP_FPS)
fourcc = cv2.VideoWriter_fourcc(*'XVID')
ret, frame = cap.read()
width = frame.shape[1]
height = frame.shape[0]


base_name = os.path.splitext(os.path.basename(video_path))[0]
output_filename = 'inferred_' + base_name + '.mp4'
out = cv2.VideoWriter(output_filename, fourcc, fps, (width, height))


while True:
    pose_model, person_model = load_combined_models()
    ret, frame = cap.read()
    # frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    # draw circle
    cv2.circle(frame, center_point, radius, (0, 0, 255), 2)
    # Detect person
    person_threshold = 0.15
    person_results = inference_detector(person_model, frame)
    person_bboxes, person_labels, _ = mmdet3x_convert_to_bboxes_mmdet(person_results, person_threshold, prefix='class')
    # remove all other class, remain class 1
    chosen_bboxes = []
    chosen_labels = []
    for bbox, label in zip(person_bboxes, person_labels):
        if label == 'class_1':
            chosen_bboxes.append(bbox)
            chosen_labels.append(label)

    # Keypoint detection
    for bbox, label in zip(chosen_bboxes, chosen_labels):
        left, top, right, bottom = bbox
        cropped_object = frame[top:bottom, left:right]
        # draw keypoint
        cropped_results = inference_topdown(pose_model, cropped_object)
        keypoints = cropped_results[0].pred_instances.keypoints
        keypoints_confidence = cropped_results[0].pred_instances.keypoint_scores
        # visualize keypoint
        vis_keypoint = visualize_all_keypoints(cropped_object)
        # plt.imshow(vis_keypoint)
        # plt.show()
        # visualize keypoint on original image
        # vis_keypoint = cv2.cvtColor(vis_keypoint, cv2.COLOR_BGR2RGB)
        frame[top:bottom, left:right] = vis_keypoint
    # draw line between center point and person
    for bbox, label in zip(chosen_bboxes, chosen_labels):
        left, top, right, bottom = bbox
        center_person = (int((left+right)/2), int((top+bottom)/2))
        cv2.line(frame, center_point, center_person, (0, 255, 0), 2)
        # show length of line
        length = math.sqrt((center_point[0]-center_person[0])**2 + (center_point[1]-center_person[1])**2)
        cv2.putText(frame, str(round(length, 2)), (center_person[0], center_person[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        if length <= trash_threshold:
            # draw red box
            cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), 2)
            # draw text
            cv2.putText(frame, 'Warning', (left, top-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
            # draw text in the frame to show warning
            cv2.putText(frame, 'Trash dumping!!!', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
    out.write(frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break
    # time.sleep(0.1)
cap.release()