''' Camera position estimation from center marker '''

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

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from default_config import *



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}")

# convert rotation matrix to euler angles (degrees)
def rotationMatrixEulerAngles(rotMatrix):
    assert rotMatrix.shape == (3, 3), "Ma trận quay phải có kích thước 3x3"
    sy = math.sqrt(rotMatrix[0, 0]* rotMatrix[0, 0] + rotMatrix[1, 0]*rotMatrix[1, 0])
    singular = sy < 1e-6
    if not singular:
        x = math.atan2(rotMatrix[2, 1], rotMatrix[2, 2])
        x = np.round(np.degrees(x))
        y = math.atan2(-rotMatrix[2, 0], sy)
        y = np.round(np.degrees(y))
        z = math.atan2(rotMatrix[1, 0], rotMatrix[0, 0])
        z = np.round(np.degrees(z))
    else:
        x = math.atan2(rotMatrix[1, 0], rotMatrix[0, 0])
        x = np.round(np.degrees(x))
        y = math.atan2(-rotMatrix[2, 0], sy)
        y = np.round(np.degrees(y))
        z = 0
    return np.array([x, y, z])

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

    if not ret:
        break
    gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    marker_corners, marker_IDs, reject = aruco.detectMarkers(
        gray_frame, marker_dict, parameters=param_markers
    )
    if marker_corners:
        rVec, tVec, _ = aruco.estimatePoseSingleMarkers(
            marker_corners, MARKER_SIZE, cam_mat, dist_coef
        )
        total_markers = range(0, marker_IDs.size)
        markers = {}
        for ids, corners, i in zip(marker_IDs, marker_corners, total_markers):
            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()
    
            # Chuyển đổi rVec thành ma trận quay R
            R, _ = cv2.Rodrigues(rVec[i])
        
            # Tính toán tọa độ các điểm của marker so với camera
            object_points = np.array([
                [-MARKER_SIZE / 2, MARKER_SIZE / 2, 0],
                [MARKER_SIZE / 2, MARKER_SIZE / 2, 0],
                [MARKER_SIZE / 2, -MARKER_SIZE / 2, 0],
                [-MARKER_SIZE / 2, -MARKER_SIZE / 2, 0]
            ])
            marker_points_camera = R @ object_points.T + tVec[i].reshape(3, 1)
            marker_points_camera = marker_points_camera.T

            distance = np.sqrt(
                tVec[i][0][2] ** 2 + tVec[i][0][0] ** 2 + tVec[i][0][1] ** 2
            )
            # Draw the pose of the marker
            # point = cv2.drawFrameAxes(frame, cam_mat, dist_coef, rVec[i], tVec[i], 4, 4)
            x = -tVec[i][0][0]
            y = tVec[i][0][1]
            z = tVec[i][0][2]
            pos = np.array([x, y, z])
            angles = rotationMatrixEulerAngles(R.T)
            if ids == 2:
                markers['marker 2'] = np.hstack((pos,angles))
            if ids == 0:
                pos += np.array([-12.3, 9.1, 0])
                markers['marker 0'] = np.hstack((pos,angles))
            if ids == 1:
                pos += np.array([12.3, 9.1, 0.])
                markers['marker 1'] = np.hstack((pos,angles))
            if ids == 3:
                pos += np.array([-7.2, -9.1, 0.])
                markers['marker 3'] = np.hstack((pos,angles))
            if ids == 4:
                pos += np.array([12.3, -5., 0.])
                markers['marker 4'] = np.hstack((pos,angles))
        if 'marker 2' in markers.keys():
            x, y, z, roll, pitch, yaw = np.round(markers['marker 2'], 1)
        else:
            cam_position = np.array([0., 0., 0., 0., 0., 0.])
            for v in markers.values():
                cam_position += v
            cam_position /= len(markers)
            x, y, z, roll, pitch, yaw = np.round(cam_position, 1)
        print(x, y, z, roll, pitch, yaw)
    # cv2.imshow("frame", frame)
    key = cv2.waitKey(1)
    if key == ord("q"):
        break
cap.release()
cv2.destroyAllWindows()