# -- coding: utf-8 --
"""
激光器-相机联合标定程序
使用方格标定板 + 激光点，拍摄多组图像，拟合相机坐标系下的激光直线方程。

流程：
  1. 加载相机内参 (calibration_result.json)
  2. 实时预览，棋盘格 + 激光点同时可见时按 y 拍照
  3. 每张照片：PnP 解算标定板位姿 → 提取激光点像素 → 反投影求交 → 得到相机坐标系 3D 点
  4. 收集所有 3D 点，PCA 拟合直线
  5. 输出：方向向量 d_cam + 基准起点 O_cam (与相机 Z=0 平面的交点)
"""

import sys
import platform
import os
from ctypes import *
import numpy as np
import cv2
import json

# ===== 标定板参数（与 calibrate.py 保持一致） =====
SQUARE_SIZE = 15.0       # 方格边长 (mm)
CHESSBOARD_COLS = 11     # 内角点列数
CHESSBOARD_ROWS = 8      # 内角点行数
DETECT_INTERVAL = 1      # 每帧都检测棋盘格
# ================================================

# ===== 激光点检测参数（蓝紫色激光，HSV 色彩空间） =====
LASER_HSV_LOWER = (100, 50, 50)   # H, S, V 下限
LASER_HSV_UPPER = (160, 255, 255) # H, S, V 上限
# ===================================================

currentsystem = platform.system()
print(f"current system: {currentsystem}")
if currentsystem == 'Windows':
    sys.path.append(os.getenv('MVCAM_COMMON_RUNENV') + "/Samples/Python/MvImport")
else:
    sys.path.append("/opt/MVS/Samples/aarch64/Python/MvImport")

from MvCameraControl_class import *

if sys.version_info[0] < 3:
    input_func = raw_input
else:
    input_func = input

# ===== 全局状态 =====
g_capture_flag = False
g_laser_points_3d = []        # 收集的相机坐标系 3D 激光点
g_image_dir = "laser_calib_images"
g_frame_count = 0
g_cb_detected = False
g_cb_corners_display = None
g_laser_display = None        # 激光点像素坐标（显示用）
g_K = None
g_dist = None

stFrame = POINTER(MV_FRAME_OUT)
fun_ctype = get_platform_functype()
FrameInfoCallBack2 = fun_ctype(None, stFrame, c_void_p, c_bool)


def IsHBPixelFormat(enPixelType=0):
    if enPixelType in (PixelType_Gvsp_HB_Mono8, PixelType_Gvsp_HB_Mono10,
                       PixelType_Gvsp_HB_Mono10_Packed, PixelType_Gvsp_HB_Mono12,
                       PixelType_Gvsp_HB_Mono12_Packed, PixelType_Gvsp_HB_Mono16,
                       PixelType_Gvsp_HB_RGB8_Packed, PixelType_Gvsp_HB_BGR8_Packed,
                       PixelType_Gvsp_HB_RGBA8_Packed, PixelType_Gvsp_HB_BGRA8_Packed,
                       PixelType_Gvsp_HB_RGB16_Packed, PixelType_Gvsp_HB_BGR16_Packed,
                       PixelType_Gvsp_HB_RGBA16_Packed, PixelType_Gvsp_HB_BGRA16_Packed,
                       PixelType_Gvsp_HB_YUV422_Packed,
                       PixelType_Gvsp_HB_YUV422_YUYV_Packed,
                       PixelType_Gvsp_HB_BayerGR8, PixelType_Gvsp_HB_BayerRG8,
                       PixelType_Gvsp_HB_BayerGB8, PixelType_Gvsp_HB_BayerBG8,
                       PixelType_Gvsp_HB_BayerRBGG8,
                       PixelType_Gvsp_HB_BayerGB10, PixelType_Gvsp_HB_BayerGB10_Packed,
                       PixelType_Gvsp_HB_BayerBG10, PixelType_Gvsp_HB_BayerBG10_Packed,
                       PixelType_Gvsp_HB_BayerRG10, PixelType_Gvsp_HB_BayerRG10_Packed,
                       PixelType_Gvsp_HB_BayerGR10, PixelType_Gvsp_HB_BayerGR10_Packed,
                       PixelType_Gvsp_HB_BayerGB12, PixelType_Gvsp_HB_BayerGB12_Packed,
                       PixelType_Gvsp_HB_BayerBG12, PixelType_Gvsp_HB_BayerBG12_Packed,
                       PixelType_Gvsp_HB_BayerRG12, PixelType_Gvsp_HB_BayerRG12_Packed,
                       PixelType_Gvsp_HB_BayerGR12, PixelType_Gvsp_HB_BayerGR12_Packed):
        return True
    return False


def IsMonoPixelFormat(enPixelType=0):
    if enPixelType in (PixelType_Gvsp_Mono8, PixelType_Gvsp_Mono10,
                       PixelType_Gvsp_Mono10_Packed, PixelType_Gvsp_Mono12,
                       PixelType_Gvsp_Mono12_Packed, PixelType_Gvsp_Mono14,
                       PixelType_Gvsp_Mono16):
        return True
    return False


def detect_laser_spot(bgr):
    """检测蓝紫色激光点：HSV 色彩空间过滤，返回亚像素质心坐标，以及二值 mask"""
    hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
    mask = cv2.inRange(hsv, LASER_HSV_LOWER, LASER_HSV_UPPER)
    # 膨胀连接断裂区域
    mask = cv2.dilate(mask, np.ones((3, 3), np.uint8), iterations=1)
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return None, None

    largest = max(contours, key=cv2.contourArea)
    M = cv2.moments(largest)
    if M["m00"] == 0:
        return None, None
    cx = M["m10"] / M["m00"]
    cy = M["m01"] / M["m00"]
    return np.array([cx, cy], dtype=np.float32), mask


def inpaint_laser(gray, laser_mask):
    """用 inpainting 填补激光区域，避免激光遮盖棋盘格角点"""
    if laser_mask is None:
        return gray
    # 膨胀 mask 留出安全边距
    kernel = np.ones((7, 7), np.uint8)
    mask_dilated = cv2.dilate(laser_mask, kernel, iterations=2)
    return cv2.inpaint(gray, mask_dilated, 3, cv2.INPAINT_TELEA)


def backproject_to_plane(laser_px, rvec, tvec, K):
    """
    将激光点像素坐标反投影到标定板平面，返回相机坐标系下的 3D 点。

    标定板平面在相机坐标系中: n · (X - t) = 0, 其中 n = R[:, 2]
    相机射线: P(s) = s * d,  其中 d = K^{-1} * [u, v, 1]^T
    求交: s = (n · t) / (n · d)
    """
    R, _ = cv2.Rodrigues(rvec)
    n = R[:, 2]               # 标定板法向量（相机坐标系）
    t = tvec.ravel()          # 标定板原点在相机坐标系的位置

    d = np.linalg.inv(K) @ np.array([laser_px[0], laser_px[1], 1.0])
    d = d / np.linalg.norm(d)  # 归一化方向

    nd = np.dot(n, d)
    if abs(nd) < 1e-8:
        return None           # 射线与平面平行

    s = np.dot(n, t) / nd
    if s <= 0:
        return None           # 交点在相机后方

    return (s * d).astype(np.float32)


def solve_pnp_for_checkerboard(gray, K, dist):
    """检测棋盘格并 PnP 解算标定板位姿"""
    ret_cb, corners = cv2.findChessboardCorners(
        gray, (CHESSBOARD_COLS, CHESSBOARD_ROWS),
        cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_NORMALIZE_IMAGE
    )
    if not ret_cb:
        return None, None

    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    corners_sub = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)

    objp = np.zeros((CHESSBOARD_COLS * CHESSBOARD_ROWS, 3), np.float32)
    objp[:, :2] = np.mgrid[0:CHESSBOARD_COLS, 0:CHESSBOARD_ROWS].T.reshape(-1, 2) * SQUARE_SIZE

    success, rvec, tvec = cv2.solvePnP(objp, corners_sub, K, dist, flags=cv2.SOLVEPNP_IPPE)
    if not success:
        return None, None

    return rvec, tvec


def fit_3d_line(points):
    """PCA 拟合空间直线，返回 (direction, centroid)"""
    pts = np.array(points)
    centroid = pts.mean(axis=0)
    centered = pts - centroid
    _, _, vh = np.linalg.svd(centered)
    direction = vh[0]  # 第一主成分 = 直线方向
    # 确保方向指向 Z+ 半球
    if direction[2] < 0:
        direction = -direction
    return direction, centroid


def image_callback(pstFrame, pUser, bAutoFree):
    global g_capture_flag, g_laser_points_3d, g_frame_count
    global g_cb_detected, g_cb_corners_display, g_laser_display
    global g_K, g_dist

    stFrame = cast(pstFrame, POINTER(MV_FRAME_OUT)).contents
    User = ctypes.cast(pUser, ctypes.py_object).value
    if not stFrame:
        return

    while True:
        stDecodeParam = MV_CC_HB_DECODE_PARAM()
        stConvertParam = MV_CC_PIXEL_CONVERT_PARAM_EX()
        memset(byref(stConvertParam), 0, sizeof(stConvertParam))

        result = IsHBPixelFormat(stFrame.stFrameInfo.enPixelType)
        if True == result:
            DecodeBufferlen = stFrame.stFrameInfo.nWidth * stFrame.stFrameInfo.nHeight * 3
            DecodeBuffer = (c_ubyte * DecodeBufferlen)()
            stDecodeParam.pSrcBuf = stFrame.pBufAddr
            stDecodeParam.nSrcLen = stFrame.stFrameInfo.nFrameLen
            stDecodeParam.pDstBuf = DecodeBuffer
            stDecodeParam.nDstBufSize = DecodeBufferlen
            ret = User.MV_CC_HBDecode(stDecodeParam)
            if ret != 0:
                print("HB Decode fail! ret[0x%x]" % ret)
                break
            stConvertParam.pSrcData = stDecodeParam.pDstBuf
            stConvertParam.nSrcDataLen = stDecodeParam.nDstBufLen
            stConvertParam.enSrcPixelType = stDecodeParam.enDstPixelType
        else:
            stConvertParam.pSrcData = stFrame.pBufAddr
            stConvertParam.nSrcDataLen = stFrame.stFrameInfo.nFrameLen
            stConvertParam.enSrcPixelType = stFrame.stFrameInfo.enPixelType

        enDstPixelType = PixelType_Gvsp_Undefined
        nChannelNum = 1
        bMono = IsMonoPixelFormat(stConvertParam.enSrcPixelType)
        if True == bMono:
            enDstPixelType = PixelType_Gvsp_Mono8
            nChannelNum = 1
        else:
            enDstPixelType = PixelType_Gvsp_BGR8_Packed
            nChannelNum = 3

        convertdestBufflen = nChannelNum * stFrame.stFrameInfo.nWidth * stFrame.stFrameInfo.nHeight
        DstBuffer = (c_ubyte * convertdestBufflen)()
        DstBufferlen = convertdestBufflen

        stConvertParam.nWidth = stFrame.stFrameInfo.nWidth
        stConvertParam.nHeight = stFrame.stFrameInfo.nHeight
        stConvertParam.enDstPixelType = enDstPixelType
        stConvertParam.pDstBuffer = DstBuffer
        stConvertParam.nDstBufferSize = DstBufferlen

        ret = User.MV_CC_ConvertPixelTypeEx(stConvertParam)
        if ret != 0:
            print("convert pixel fail! ret[0x%x]" % ret)
            break

        if nChannelNum == 1:
            numpy_image = np.frombuffer(DstBuffer, dtype=np.ubyte, count=DstBufferlen).reshape(
                stFrame.stFrameInfo.nHeight, stFrame.stFrameInfo.nWidth)
            gray_full = numpy_image
            display_image = cv2.cvtColor(numpy_image, cv2.COLOR_GRAY2BGR)
        else:
            numpy_image = np.frombuffer(DstBuffer, dtype=np.ubyte, count=DstBufferlen).reshape(
                stFrame.stFrameInfo.nHeight, stFrame.stFrameInfo.nWidth, 3)
            gray_full = cv2.cvtColor(numpy_image, cv2.COLOR_BGR2GRAY)
            display_image = cv2.resize(numpy_image, None, fx=0.5, fy=0.5, interpolation=cv2.INTER_AREA)

        # 激光点检测（彩色图，HSV 色彩空间）
        if nChannelNum == 3:
            laser_px, laser_mask = detect_laser_spot(numpy_image)
        else:
            laser_px, laser_mask = None, None

        # 用 inpainting 填补激光区域后再做棋盘格检测，避免激光遮盖角点
        if laser_mask is not None:
            gray_inpainted = inpaint_laser(gray_full, laser_mask)
        else:
            gray_inpainted = gray_full

        if g_frame_count % DETECT_INTERVAL == 0:
            gray_disp = cv2.resize(gray_inpainted, None, fx=0.5, fy=0.5, interpolation=cv2.INTER_AREA) if nChannelNum == 3 else gray_inpainted
            ret_cb, corners_disp = cv2.findChessboardCorners(
                gray_disp, (CHESSBOARD_COLS, CHESSBOARD_ROWS),
                cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_NORMALIZE_IMAGE
            )
            g_cb_detected = ret_cb
            if ret_cb:
                criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
                g_cb_corners_display = cv2.cornerSubPix(gray_disp, corners_disp, (5, 5), (-1, -1), criteria)
            else:
                g_cb_corners_display = None

        if laser_px is not None and nChannelNum == 3:
            g_laser_display = laser_px * 0.5
        elif laser_px is not None:
            g_laser_display = laser_px.copy()
        else:
            g_laser_display = None

        # 拍照
        if g_capture_flag:
            if g_cb_detected and laser_px is not None and g_K is not None:
                # PnP 使用 inpainting 后的灰度图，避免激光干扰角点提取
                rvec, tvec = solve_pnp_for_checkerboard(gray_inpainted, g_K, g_dist)
                if rvec is not None:
                    p3d = backproject_to_plane(laser_px, rvec, tvec, g_K)
                    if p3d is not None:
                        g_laser_points_3d.append(p3d)
                        idx = len(g_laser_points_3d)
                        path = os.path.join(g_image_dir, f"laser_{idx:02d}.png")
                        cv2.imwrite(path, numpy_image)
                        print(f"  [OK] 第 {idx} 个激光点: "
                              f"像素({laser_px[0]:.1f}, {laser_px[1]:.1f}) -> "
                              f"3D({p3d[0]:.1f}, {p3d[1]:.1f}, {p3d[2]:.1f}) mm")
                        print(f"       图像已保存: {path}")
                        g_capture_flag = False
                    else:
                        print("  警告: 激光点反投影失败")
                        g_capture_flag = False
                else:
                    print("  警告: PnP 解算失败")
                    g_capture_flag = False
            else:
                pass  # 等待检测到棋盘格和激光点

        # 绘制
        if g_cb_detected and g_cb_corners_display is not None:
            cv2.drawChessboardCorners(display_image, (CHESSBOARD_COLS, CHESSBOARD_ROWS),
                                      g_cb_corners_display, True)
            cv2.putText(display_image, "CHESSBOARD DETECTED", (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        if g_laser_display is not None:
            lx, ly = int(g_laser_display[0]), int(g_laser_display[1])
            cv2.circle(display_image, (lx, ly), 2, (0, 0, 255), -1)
            cv2.circle(display_image, (lx, ly), 12, (0, 0, 255), 2)
            cv2.putText(display_image, "LASER", (lx + 15, ly - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)

        status = f"Points: {len(g_laser_points_3d)}"
        cv2.putText(display_image, status, (10, display_image.shape[0] - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)

        cv2.imshow('Laser Calibration', display_image)
        cv2.waitKey(1)

        g_frame_count += 1
        break

    if bAutoFree == False:
        if pUser != None:
            User.MV_CC_FreeImageBuffer(stFrame)


CALL_BACK_FUN = FrameInfoCallBack2(image_callback)


def run_fitting():
    """拟合激光直线并输出结果"""
    if len(g_laser_points_3d) < 3:
        print(f"错误：激光点数量不足（需要至少3个，当前{len(g_laser_points_3d)}个）！")
        return None, None

    direction, centroid = fit_3d_line(g_laser_points_3d)

    # 计算与相机 Z=0 平面的交点作为参考原点
    if abs(direction[2]) > 1e-8:
        t_z0 = -centroid[2] / direction[2]
        origin = centroid + t_z0 * direction
    else:
        origin = centroid.copy()
        print("警告：激光方向近似平行于相机 Z=0 平面，使用质心作为参考原点")

    # 计算拟合残差
    centered = np.array(g_laser_points_3d) - centroid
    residuals = np.linalg.norm(centered - np.outer(np.dot(centered, direction), direction), axis=1)
    rms = np.sqrt(np.mean(residuals ** 2))

    print("\n" + "=" * 55)
    print("           激光直线标定结果")
    print("=" * 55)
    print(f"\n  采集点数: {len(g_laser_points_3d)}")
    print(f"  RMS 拟合残差: {rms:.3f} mm")
    print(f"\n  直线参数 (相机坐标系):")
    print(f"  方向向量 d_cam:  [{direction[0]:.6f}, {direction[1]:.6f}, {direction[2]:.6f}]")
    print(f"  基准起点 O_cam:  [{origin[0]:.2f}, {origin[1]:.2f}, {origin[2]:.2f}] mm")
    print(f"  (O_cam 为直线与相机 Z=0 平面的交点)")
    print("=" * 55)

    # 保存
    result = {
        "direction": direction.tolist(),
        "origin": origin.tolist(),
        "num_points": len(g_laser_points_3d),
        "rms_error": float(rms),
        "points_3d": [p.tolist() for p in g_laser_points_3d],
    }
    save_path = "laser_result.json"
    with open(save_path, "w") as f:
        json.dump(result, f, indent=2)
    print(f"\n结果已保存到: {save_path}")

    # 打印可粘贴到 laser_pnp.py 的 JSON
    paste = {"direction": direction.tolist(), "origin": origin.tolist()}
    print(f"\n--- 复制以下 JSON 到 laser_pnp.py 的 LASER_JSON 变量中 ---")
    print(json.dumps(paste, indent=2))
    print(f"---------------------------------------------------------")

    return direction, origin


if __name__ == "__main__":
    # 加载内参
    calib_path = "calibration_result.json"
    try:
        with open(calib_path, "r") as f:
            calib_data = json.load(f)
        g_K = np.array(calib_data["mtx"])
        g_dist = np.array(calib_data["dist"])
        img_size = tuple(calib_data["img_size"])
        print(f"已加载内参: {calib_path}")
        print(f"  分辨率: {img_size[0]} x {img_size[1]}")
    except FileNotFoundError:
        print(f"未找到标定文件: {calib_path}")
        print("请先运行 calibrate.py 进行相机内参标定")
        sys.exit(1)

    print("=" * 55)
    print("       激光器-相机联合标定程序")
    print(f"       棋盘格: {CHESSBOARD_COLS}x{CHESSBOARD_ROWS}")
    print(f"       方格边长: {SQUARE_SIZE} mm")
    print("=" * 55)
    print()

    # 输入拍照数量
    while True:
        try:
            n_photos = int(input_func("请输入要拍摄的图像数量 (建议 >= 10): "))
            if n_photos >= 3:
                break
            print("至少需要3张！")
        except ValueError:
            print("请输入有效数字！")

    print(f"\n将拍摄 {n_photos} 张图像。")
    print("操作说明：")
    print("  1. 将标定板以不同姿态放在相机视野中")
    print("  2. 确保激光点投射在标定板上")
    print("  3. 看到绿色棋盘格 + 红色激光点标注后，输入 y 并回车")
    print()

    # ===== 初始化相机 =====
    MvCamera.MV_CC_Initialize()

    deviceList = MV_CC_DEVICE_INFO_LIST()
    tlayerType = (MV_GIGE_DEVICE | MV_USB_DEVICE | MV_GENTL_CAMERALINK_DEVICE
                  | MV_GENTL_CXP_DEVICE | MV_GENTL_XOF_DEVICE)

    ret = MvCamera.MV_CC_EnumDevices(tlayerType, deviceList)
    if ret != 0:
        print("enum devices fail! ret[0x%x]" % ret)
        sys.exit()

    if deviceList.nDeviceNum == 0:
        print("find no device!")
        sys.exit()

    print("find %d device(s)" % deviceList.nDeviceNum)

    nConnectionNum = 0
    cam = MvCamera()
    stDeviceList = cast(deviceList.pDeviceInfo[int(nConnectionNum)], POINTER(MV_CC_DEVICE_INFO)).contents

    ret = cam.MV_CC_CreateHandle(stDeviceList)
    if ret != 0:
        print("create handle fail! ret[0x%x]" % ret)
        sys.exit()

    ret = cam.MV_CC_OpenDevice(MV_ACCESS_Exclusive, 0)
    if ret != 0:
        print("open device fail! ret[0x%x]" % ret)
        sys.exit()

    if stDeviceList.nTLayerType == MV_GIGE_DEVICE or stDeviceList.nTLayerType == MV_GENTL_GIGE_DEVICE:
        nPacketSize = cam.MV_CC_GetOptimalPacketSize()
        if int(nPacketSize) > 0:
            cam.MV_CC_SetIntValue("GevSCPSPacketSize", nPacketSize)

    cam.MV_CC_SetEnumValue("TriggerMode", MV_TRIGGER_MODE_OFF)

    ret = cam.MV_CC_RegisterImageCallBackEx2(CALL_BACK_FUN, ctypes.py_object(cam), True)
    if ret != 0:
        print("register image callback fail! ret[0x%x]" % ret)
        sys.exit()

    ret = cam.MV_CC_StartGrabbing()
    if ret != 0:
        print("start grabbing fail! ret[0x%x]" % ret)
        sys.exit()

    print("相机取流已启动。\n")

    os.makedirs(g_image_dir, exist_ok=True)
    print(f"标定图像将保存到: {g_image_dir}/\n")

    # ===== 交互式拍照 =====
    try:
        while len(g_laser_points_3d) < n_photos:
            remaining = n_photos - len(g_laser_points_3d)
            cmd = input_func(f"[{len(g_laser_points_3d)}/{n_photos}] 输入 y 并回车拍照 (剩余 {remaining}): ")

            if cmd.strip().lower() == 'y':
                g_capture_flag = True
                waited = 0
                while g_capture_flag and waited < 100:
                    cv2.waitKey(50)
                    waited += 1
                if g_capture_flag:
                    print("  警告: 超时，请确保棋盘格和激光点同时可见")
                    g_capture_flag = False
            elif cmd.strip().lower() == 'q':
                print("用户取消。")
                break
            else:
                print("  请输入 y (拍照) 或 q (退出)")
    except KeyboardInterrupt:
        print("\n用户中断。")

    # ===== 停止取流 =====
    print("\n停止相机取流...")
    cv2.destroyAllWindows()

    cam.MV_CC_StopGrabbing()
    cam.MV_CC_CloseDevice()
    cam.MV_CC_DestroyHandle()
    MvCamera.MV_CC_Finalize()

    # ===== 拟合激光直线 =====
    if len(g_laser_points_3d) > 0:
        run_fitting()
    else:
        print("\n没有采集到激光点，跳过拟合。")