import sys
import numpy as np
import cv2
import cv2.aruco as aruco
import math
from ctypes import *

# MVS 파이썬 SDK 경로 추가
sys.path.append("/opt/MVS/Samples/64/Python/MvImport")
from MvCameraControl_class import *

def rotation_matrix_to_quaternion(R):
    """
    회전 행렬(Rotation Matrix)을 수치적으로 안정적인 쿼터니언(qx, qy, qz, qw)으로 변환
    """
    tr = R[0, 0] + R[1, 1] + R[2, 2]

    if tr > 0:
        S = math.sqrt(tr + 1.0) * 2
        qw = 0.25 * S
        qx = (R[2, 1] - R[1, 2]) / S
        qy = (R[0, 2] - R[2, 0]) / S
        qz = (R[1, 0] - R[0, 1]) / S
    elif (R[0, 0] > R[1, 1]) and (R[0, 0] > R[2, 2]):
        S = math.sqrt(1.0 + R[0, 0] - R[1, 1] - R[2, 2]) * 2
        qw = (R[2, 1] - R[1, 2]) / S
        qx = 0.25 * S
        qy = (R[0, 1] + R[1, 0]) / S
        qz = (R[0, 2] + R[2, 0]) / S
    elif R[1, 1] > R[2, 2]:
        S = math.sqrt(1.0 + R[1, 1] - R[0, 0] - R[2, 2]) * 2
        qw = (R[0, 2] - R[2, 0]) / S
        qx = (R[0, 1] + R[1, 0]) / S
        qy = 0.25 * S
        qz = (R[1, 2] + R[2, 1]) / S
    else:
        S = math.sqrt(1.0 + R[2, 2] - R[0, 0] - R[1, 1]) * 2
        qw = (R[1, 0] - R[0, 1]) / S
        qx = (R[0, 2] + R[2, 0]) / S
        qy = (R[1, 2] + R[2, 1]) / S
        qz = 0.25 * S
        
    if tr <= 0:
        if (R[0, 0] > R[1, 1]) and (R[0, 0] > R[2, 2]): pass
        elif R[1, 1] > R[2, 2]: qw = (R[0, 2] - R[2, 0]) / S
        else: qw = (R[1, 0] - R[0, 1]) / S

    return qx, qy, qz, qw

def main():
    # 1. 카메라 초기화
    deviceList = MV_CC_DEVICE_INFO_LIST()
    tlayerType = MV_GIGE_DEVICE | MV_USB_DEVICE
    ret = MvCamera.MV_CC_EnumDevices(tlayerType, deviceList)
    if ret != 0:
        print(f"카메라 검색 실패! ret[0x{ret:x}]")
        return

    if deviceList.nDeviceNum == 0:
        print("연결된 카메라가 없습니다.")
        return

    stDeviceImgLimit = cast(deviceList.pDeviceInfo[0], POINTER(MV_CC_DEVICE_INFO)).contents
    cam = MvCamera()
    cam.MV_CC_CreateHandle(stDeviceImgLimit)
    cam.MV_CC_OpenDevice(MV_ACCESS_Exclusive, 0)
    
    cam.MV_CC_SetEnumValue("ExposureAuto", 0) 
    cam.MV_CC_SetFloatValue("ExposureTime", 50000.0) 
    cam.MV_CC_SetBoolValue("AcquisitionFrameRateEnable", True)
    cam.MV_CC_SetFloatValue("AcquisitionFrameRate", 15.0) 

    ret = cam.MV_CC_StartGrabbing()
    
    # 150mm 마커 크기 반영 (0.15미터)
    marker_size = 0.15
    
    # 캘리브레이션 파라미터 (1440x1080 규격 원본 파라미터)
    fx, fy, cx, cy = 1203.078148, 1206.096396, 699.186863, 565.715472
    mtx = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
    dist = np.array([[-0.102740, 0.093985, -0.000759, -0.001804, 0.0]])
    
    # 7x7 출력 마커 설정
    dictionary = aruco.getPredefinedDictionary(aruco.DICT_7X7_50)
    parameters = aruco.DetectorParameters()

    try:
        while True:
            stOutFrame = MV_FRAME_OUT()
            ret = cam.MV_CC_GetImageBuffer(stOutFrame, 1000)
            if ret == 0:
                if stOutFrame.pBufAddr:
                    addr = cast(stOutFrame.pBufAddr, c_void_p).value
                    nWidth = stOutFrame.stFrameInfo.nWidth
                    nHeight = stOutFrame.stFrameInfo.nHeight
                    nFrameLen = stOutFrame.stFrameInfo.nFrameLen
                    
                    pData = (c_ubyte * nFrameLen).from_address(addr)
                    img_raw = np.frombuffer(pData, dtype=np.uint8).reshape(nHeight, nWidth)
                    frame = cv2.cvtColor(img_raw, cv2.COLOR_BayerGB2BGR)

                    # ArUco 인식
                    corners, ids, _ = aruco.detectMarkers(frame, dictionary, parameters=parameters)
                    if ids is not None:
                        # 꼭짓점 좌표 정밀도를 소수점 이하(Sub-pixel) 단위로 보정하는 연산
                        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.001)
                        for corner in corners:
                            cv2.cornerSubPix(gray, corner, (5, 5), (-1, -1), criteria)

                        for i in range(len(ids)):
                            marker_id = ids[i][0]
                            rvec, tvec, _ = aruco.estimatePoseSingleMarkers(corners[i], marker_size, mtx, dist)
                            
                            # 회전 변환 행렬 계산
                            rmat, _ = cv2.Rodrigues(rvec)
                            
                            # 💡 [수정] OpenCV 기준계와 매칭하기 위한 X축 180도 회전 보정 행렬
                            R_flip = np.array([
                                [1,  0,  0],
                                [0, -1,  0],
                                [0,  0, -1]
                            ], dtype=np.float32)
                            
                            # 회전 행렬 교정 연산
                            rmat_corrected = np.dot(rmat, R_flip)
                            
                            # 보정된 회전 행렬 기반으로 쿼터니언 변환 계산
                            qx, qy, qz, qw = rotation_matrix_to_quaternion(rmat_corrected)
                            
                            # 시각화할 좌표축 벡터도 보정 처리
                            rvec_corrected, _ = cv2.Rodrigues(rmat_corrected)
                            
                            # 마커 '정중앙' 원점 추출 및 수치 맵핑
                            tvec_center = tvec[0][0]
                            x, y, z = tvec_center[0], tvec_center[1], tvec_center[2]
                            
                            # 시각화 (마커 테두리선 및 3축 기둥 렌더링)
                            aruco.drawDetectedMarkers(frame, corners)
                            cv2.drawFrameAxes(frame, mtx, dist, rvec_corrected, tvec_center, 0.05)
                            
                            # 텍스트 출력 위치 설정
                            text_pos = (int(corners[i][0][0][0]), int(corners[i][0][0][1]) - 75)
                            
                            # 텍스트 색상을 실제 시각화 축 색상과 1:1 일치 (BGR 포맷 기준)
                            cv2.putText(frame, f"ID:{marker_id}", text_pos, cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255), 2)
                            cv2.putText(frame, f"X:{x:.2f}m (Red)", (text_pos[0], text_pos[1] + 20), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 255), 2)
                            cv2.putText(frame, f"Y:{y:.2f}m (Green)", (text_pos[0], text_pos[1] + 40), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 255, 0), 2)
                            cv2.putText(frame, f"Z:{z:.2f}m (Blue)", (text_pos[0], text_pos[1] + 60), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 0, 0), 2)
                            
                            # 수치 안정성이 확보된 쿼터니언 정보 출력
                            quat_text = f"Q [x:{qx:.3f}, y:{qy:.3f}, z:{qz:.3f}, w:{qw:.3f}]"
                            cv2.putText(frame, quat_text, (text_pos[0], text_pos[1] + 80), cv2.FONT_HERSHEY_SIMPLEX, 0.42, (0, 255, 255), 2)

                    cv2.imshow("Hikrobot ArUco V2 (Pose with Units)", frame)
                
                cam.MV_CC_FreeImageBuffer(stOutFrame)
            
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
    finally:
        cam.MV_CC_StopGrabbing()
        cam.MV_CC_CloseDevice()
        cam.MV_CC_DestroyHandle()
        cv2.destroyAllWindows()

if __name__ == "__main__":
    main()