''' 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 *

visualization = True

# Tải dữ liệu hiệu chuẩn camera, phù hợp với độ phân giải ban đầu
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()
    center_y, center_x = frame.shape[0]//2, frame.shape[1]//2
    cv2.drawMarker(frame, (center_x, center_y), (0,0,255), markerType=cv2.MARKER_CROSS, markerSize=10, thickness=1, line_type=cv2.LINE_AA)
    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.linalg.norm(tVec[i][0])

            # Draw the pose of the marker
            point = cv2.drawFrameAxes(frame, cam_mat, dist_coef, rVec[i], tVec[i], 4, 4)


            # R, _ = cv2.Rodrigues(rVec[i])        # Chuyển rVec sang ma trận quay
            # R_inv = R.T           # Lấy nghịch đảo của ma trận quay (tức là R^T nếu R là trực giao)
            # position = - R_inv @ tVec[i].reshape(3, 1)  # Tính vị trí camera trong hệ marker
            # 
            # 
            # position = position.reshape(3,)
            
            position =np.array([-tVec[i][0][0], tVec[i][0][1], tVec[i][0][2]])
            position += default_positions[ids[0]]

            angles = rotationMatrixEulerAngles(R.T)
            markers[f'marker {ids[0]}'] = np.hstack((position, angles))

            if visualization:
                cv2.putText(
                    frame,
                    f"id: {ids[0]}",
                    top_right,
                    cv2.FONT_HERSHEY_PLAIN,
                    1.3,
                    (0, 0, 255), 2, cv2.LINE_AA
                )
        # Nếu marker trung tâm (marker 0) xuất hiện thì chỉ cần lấy vị trí so với marker này
        if 'marker 0' in markers.keys():
            cam_position = markers['marker 0']
        # Nếu không, xét trung bình các marker
        else:
            cam_position = np.array([0., 0., 0., 0., 0., 0.])
            for v in markers.values():
                cam_position += v
            cam_position /= len(markers)

        print(','.join(f'{i:.1f}' for i in cam_position.tolist()))
        if visualization:
            cv2.putText(
                frame,
                ','.join(f'{i:.1f}' for i in cam_position[:3].tolist()),
                (10, 20),
                cv2.FONT_HERSHEY_PLAIN,
                1.,
                (0, 255, 0), 1, cv2.LINE_AA
            )
            cv2.putText(
                frame,
                ','.join(f'{i:.1f}' for i in cam_position[3:].tolist()),
                (10, 40),
                cv2.FONT_HERSHEY_PLAIN,
                1.,
                (0, 0, 255), 1, cv2.LINE_AA
            )

    if visualization:
        cv2.imshow("frame", frame)
        key = cv2.waitKey(1)
        if key == ord("q"):
            break
cap.release()

if visualization:
    cv2.destroyAllWindows()