import cv2
from cv2 import aruco
import numpy as np
import os
import sys

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from default_config import *

visualization = True


calib_data_path = f"camera_calibration/calib_data/{set_resolution}/MultiMatrix.npz"
calib_data = np.load(calib_data_path)
print('resolution:', set_resolution)

cam_mat = calib_data["camMatrix"]
dist_coef = calib_data["distCoef"]


marker_dict = aruco.getPredefinedDictionary(aruco.DICT_7X7_50)
param_markers = aruco.DetectorParameters()

cap = cv2.VideoCapture(0)

# Thay đổi độ phân giải phù hợp với yêu cầu ban đầu
new_height = RESOLUTION_STANDARDS[set_resolution]
new_width, new_height = change_resolution(cap,new_height)    
print(f'Resolution camera {new_width}x{new_height}')

fps = cap.get(cv2.CAP_PROP_FPS)
print(f"Frames per second: {fps}")

while True:
    ret, frame = cap.read()

    if not ret:
        break
    gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    detector = cv2.aruco.ArucoDetector(marker_dict, param_markers)
    marker_corners, marker_IDs, reject = detector.detectMarkers(gray_frame)
    if marker_corners:
        rVec, tVec, _ = aruco.estimatePoseSingleMarkers(
            marker_corners, MARKER_SIZE, cam_mat, dist_coef
        )
        total_markers = range(0, marker_IDs.size)
        for ids, corners, i in zip(marker_IDs, marker_corners, total_markers):
            if visualization:
                cv2.polylines(
                    frame, [corners.astype(np.int32)], True, (0, 255, 255), 4, cv2.LINE_AA
                )
            corners = corners.reshape(4, 2)
            corners = corners.astype(int)
            top_right = corners[0].ravel()
            top_left = corners[1].ravel()
            bottom_right = corners[2].ravel()
            bottom_left = corners[3].ravel()

            distance = np.linalg.norm(tVec[i][0])
            x = tVec[i][0][0]
            y = tVec[i][0][1]
            z = tVec[i][0][2]

            # Draw the pose of the marker
            if visualization:
                point = cv2.drawFrameAxes(frame, cam_mat, dist_coef, rVec[i], tVec[i], 4, 4)
                cv2.putText(
                    frame,
                    f"id: {ids[0]} Dist: {distance: .2f}",
                    top_right,
                    cv2.FONT_HERSHEY_PLAIN,
                    1.3,
                    (0, 0, 255),
                    2,
                    cv2.LINE_AA,
                )
                cv2.putText(
                    frame,
                    f"x:{x: .1f} y: {y: .1f} ",
                    bottom_right,
                    cv2.FONT_HERSHEY_PLAIN,
                    1.0,
                    (0, 0, 255),
                    2,
                    cv2.LINE_AA,
                )
            print(f"Marker {ids[0]}: x = {x: .1f}, y = {y: .1f} , z = {z: .1f}, distance = {distance: .2f}")
    if visualization:
        cv2.imshow("frame", frame)
        key = cv2.waitKey(1)
        if key == ord("q"):
            break
cap.release()

if visualization:
    cv2.destroyAllWindows()