from collections import defaultdict

import cv2
import numpy as np
import torch

from ultralytics import YOLO

from orbbec_cam import OrbbecCamera

class Vision():
    def __init__(self, 
                 model='yolov8n.pt',
                 color=True,
                 depth=False) -> None:
        # init YOLO model
        self.model = YOLO(model)
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model.to(device)

        # init Orbbec camera
        self.orbbec_cam = OrbbecCamera(color=color,depth=depth)


    def color_view(self):
        while True:
            # read a frame from color camera
            try:
                color_image = self.orbbec_cam.color_viewer.get_color_image()
                if color_image is None:
                    continue
                # Display the color image
                cv2.imshow("Color Viewer", color_image)

                key = cv2.waitKey(1)
                # press 'q' or ESC to exit
                if key == ord('q') or key == 27:
                    break

            except KeyboardInterrupt:
                break
        
        # release the video capture object and close the display window
        self.orbbec_cam.pipeline.stop()
        cv2.destroyAllWindows()


    def depth_view(self):
        while True:
            # read a frame from depth camera
            try:
                depth_image = self.orbbec_cam.depth_viewer.get_depth_image()
                if depth_image is None:
                    continue

                # Display the depth image
                cv2.imshow("Depth Viewer", depth_image)

                key = cv2.waitKey(1)
                # press 'q' or ESC to exit
                if key == ord('q') or key == 27:
                    break

            except KeyboardInterrupt:
                break
        
        # release the video capture object and close the display window
        self.orbbec_cam.pipeline.stop()
        cv2.destroyAllWindows()


    def track(self):
        # Store the track history
        track_history = defaultdict(lambda: [])

        while True:
            try:
                # Read a frame from orbbec color camera
                color_image = self.orbbec_cam.color_viewer.get_color_image()
                if color_image is None:
                    continue
                
                # Run tracking on the frame, persisting tracks between frames
                results = self.model.track(color_image, persist=True)

                # If no object detected, show the original color image
                if results[0].boxes.id is None:
                    annotated_frame = color_image
                else:
                    # Get the boxes and track IDs
                    boxes = results[0].boxes.xywh.cpu()
                    track_ids = results[0].boxes.id.int().cpu().tolist()

                    # Visualize the results on the frame
                    annotated_frame = results[0].plot()

                    # Plot the tracks
                    for box, track_id in zip(boxes, track_ids):
                        x, y, w, h = box
                        track = track_history[track_id]
                        track.append((float(x), float(y)))  # x, y center point
                        if len(track) > 30:  # retain 30 tracks for 30 frames
                            track.pop(0)

                        # Draw the tracking lines
                        points = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
                        cv2.polylines(annotated_frame, [points], isClosed=False, color=(255, 255, 255), thickness=2)

                        # draw the distance label
                        depth_data = self.orbbec_cam.depth_viewer.get_depth_data()
                        if depth_data is None:
                            continue

                        cord = (int(y/annotated_frame.shape[0]*depth_data.shape[0]),
                                int(x/annotated_frame.shape[1]*depth_data.shape[1]))
                        print(depth_data[cord])
                        cv2.putText(annotated_frame,
                                    str(depth_data[cord]),
                                    (int(x), int(y)),
                                    cv2.FONT_HERSHEY_SIMPLEX,
                                    0.5, (0, 0, 255), 1)

                
                # Display the annotated frame
                cv2.imshow("Tracking", annotated_frame)

                key = cv2.waitKey(1)
                # press 'q' or ESC to exit
                if key == ord('q') or key == 27:
                    break

            except KeyboardInterrupt:
                break

        # release the video capture object and close the display window
        self.orbbec_cam.pipeline.stop()
        cv2.destroyAllWindows()


if __name__=="__main__":
    vision = Vision('yolov8n-seg.pt', color=True, depth=True)
    vision.track()
    # vision.color_view()
    # vision.depth_view()