"""
AprilTag检测与像素标定模块

功能：
1. 检测图像中的AprilTag 36h11标签
2. 计算像素/毫米的转换比例
3. 提供亚像素精度的标定参数
"""

import cv2
import numpy as np
from typing import Tuple, Optional, Dict
import logging

try:
    import dt_apriltags as apriltag
    APRILTAG_AVAILABLE = True
except ImportError:
    APRILTAG_AVAILABLE = False
    logging.error("未找到可用的AprilTag库，请安装：pip install dt-apriltags")


class AprilTagCalibrator:
    """AprilTag标定器"""

    def __init__(self, tag_size_mm: float = 37.58, tag_family: str = 'tag16h5'):
        """
        初始化AprilTag标定器

        参数:
            tag_size_mm: AprilTag外部黑框边长（毫米）
            tag_family: AprilTag家族类型
        """
        if not APRILTAG_AVAILABLE:
            raise ImportError("AprilTag库未安装，请运行：pip install apriltag")

        self.tag_size_mm = tag_size_mm
        self.tag_family = tag_family

        # 配置检测器参数（dt-apriltags接口）
        self.detector = apriltag.Detector(
            families=tag_family,
            nthreads=4,
            quad_decimate=1.0,
            quad_sigma=0.0,
            refine_edges=1,
            decode_sharpening=0.25
        )
        logging.info(f"AprilTag检测器初始化完成：{tag_family}, 尺寸={tag_size_mm}mm")

    def detect(self, image: np.ndarray, use_homography: bool = True) -> Optional[Dict]:
        """
        检测图像中的AprilTag并计算标定参数（支持透视校正）

        参数:
            image: BGR格式图像
            use_homography: 是否计算单应性矩阵进行透视校正

        返回:
            标定结果字典，包含：
            - detected: 是否检测到Tag
            - tag_id: Tag的ID
            - px_per_mm: 像素/毫米比例
            - px_per_mm_corrected: 校正后的像素/毫米比例（如果use_homography=True）
            - corners: 4个角点坐标 (4x2数组)
            - center: Tag中心坐标
            - tag_size_px: Tag在图像中的平均边长（像素）
            - homography: 单应性矩阵（用于透视校正）
            - reprojection_error: 重投影误差（质量指标）

            如果未检测到返回None
        """
        # 转换为灰度图
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image

        # 检测AprilTag
        detections = self.detector.detect(gray)

        if len(detections) == 0:
            logging.warning("未检测到AprilTag")
            return None

        # 如果检测到多个Tag，选择最大的（可能最清晰）
        if len(detections) > 1:
            logging.info(f"检测到{len(detections)}个AprilTag，选择面积最大的")
            detections = sorted(detections, key=lambda d: self._compute_tag_area(d.corners), reverse=True)

        detection = detections[0]

        # 提取角点（按顺序：左下、右下、右上、左上）
        corners_detected = detection.corners.astype(np.float32)

        # 计算Tag的4条边长（像素）
        side_lengths_px = [
            np.linalg.norm(corners_detected[0] - corners_detected[1]),  # 底边
            np.linalg.norm(corners_detected[1] - corners_detected[2]),  # 右边
            np.linalg.norm(corners_detected[2] - corners_detected[3]),  # 顶边
            np.linalg.norm(corners_detected[3] - corners_detected[0])   # 左边
        ]

        # 平均边长（用于初始估计）
        avg_side_px = np.mean(side_lengths_px)
        std_side_px = np.std(side_lengths_px)

        # 质量检查：如果4条边长差异太大，说明存在透视畸变
        distortion_ratio = std_side_px / avg_side_px
        if distortion_ratio > 0.05:  # 标准差超过5%
            logging.warning(
                f"AprilTag边长不一致（std/mean={distortion_ratio:.3f}），"
                f"检测到透视畸变，建议使用单应性矩阵校正"
            )

        # 简单方法的像素/毫米比例（未校正）
        px_per_mm_simple = avg_side_px / self.tag_size_mm

        # 计算中心点
        center = detection.center

        result = {
            'detected': True,
            'tag_id': detection.tag_id,
            'px_per_mm': px_per_mm_simple,  # 未校正的比例
            'corners': corners_detected,
            'center': center,
            'tag_size_px': avg_side_px,
            'side_lengths_px': side_lengths_px,
            'side_std_px': std_side_px,
            'distortion_ratio': distortion_ratio
        }

        # ========== 单应性矩阵校正（关键优化） ==========
        if use_homography:
            homography_result = self._compute_homography(corners_detected, avg_side_px)
            result.update(homography_result)

            logging.info(
                f"AprilTag检测成功 - ID:{detection.tag_id}, "
                f"边长:{avg_side_px:.2f}px, "
                f"简单比例:{px_per_mm_simple:.3f}px/mm, "
                f"校正比例:{result.get('px_per_mm_corrected', px_per_mm_simple):.3f}px/mm, "
                f"重投影误差:{result.get('reprojection_error', 0):.3f}px"
            )
        else:
            logging.info(
                f"AprilTag检测成功 - ID:{detection.tag_id}, "
                f"边长:{avg_side_px:.2f}px, "
                f"比例:{px_per_mm_simple:.3f}px/mm"
            )

        return result

    def _compute_homography(self, corners_detected: np.ndarray, avg_side_px: float) -> Dict:
        """
        计算单应性矩阵和校正后的标定参数

        原理：
        1. AprilTag是已知的正方形
        2. 利用4个角点计算从畸变图像到理想正方形的单应性矩阵
        3. 用单应性矩阵校正后的Tag边长计算精确的px/mm比例

        参数:
            corners_detected: 检测到的4个角点（实际图像坐标，有畸变）
            avg_side_px: 平均边长（用于确定理想正方形大小）

        返回:
            包含单应性矩阵和校正参数的字典
        """
        # 构建理想正方形的4个角点（以Tag中心为原点）
        # 使用平均边长作为理想正方形的边长
        half_size = avg_side_px / 2

        # 计算检测到的Tag中心
        center_x = np.mean(corners_detected[:, 0])
        center_y = np.mean(corners_detected[:, 1])

        # 理想正方形的角点（顺序与corners_detected一致：左下、右下、右上、左上）
        corners_ideal = np.array([
            [center_x - half_size, center_y + half_size],  # 左下
            [center_x + half_size, center_y + half_size],  # 右下
            [center_x + half_size, center_y - half_size],  # 右上
            [center_x - half_size, center_y - half_size]   # 左上
        ], dtype=np.float32)

        # 计算单应性矩阵：从实际检测点 → 理想正方形
        H, mask = cv2.findHomography(corners_detected, corners_ideal, cv2.RANSAC, 5.0)

        if H is None:
            logging.error("单应性矩阵计算失败")
            return {
                'homography': None,
                'px_per_mm_corrected': avg_side_px / self.tag_size_mm,
                'reprojection_error': float('inf')
            }

        # 计算重投影误差（质量指标）
        corners_reprojected = cv2.perspectiveTransform(
            corners_detected.reshape(-1, 1, 2),
            H
        ).reshape(-1, 2)

        reprojection_error = np.mean(
            np.linalg.norm(corners_reprojected - corners_ideal, axis=1)
        )

        # 使用理想正方形的边长计算校正后的px/mm比例
        # 这个比例是在"校正后的图像"上的比例
        px_per_mm_corrected = avg_side_px / self.tag_size_mm

        # 计算缩放因子（用于评估畸变程度）
        # 如果完全无畸变，H应该接近单位矩阵
        scale_factor = np.sqrt(np.abs(np.linalg.det(H[:2, :2])))

        return {
            'homography': H,
            'homography_inv': np.linalg.inv(H),
            'px_per_mm_corrected': px_per_mm_corrected,
            'reprojection_error': reprojection_error,
            'scale_factor': scale_factor,
            'corners_ideal': corners_ideal
        }

    def visualize(self, image: np.ndarray, calibration_result: Dict) -> np.ndarray:
        """
        在图像上可视化AprilTag检测结果

        参数:
            image: 原始BGR图像
            calibration_result: detect()返回的标定结果

        返回:
            绘制了标注的图像副本
        """
        vis_image = image.copy()

        if not calibration_result or not calibration_result['detected']:
            return vis_image

        corners = calibration_result['corners']
        center = calibration_result['center']
        tag_id = calibration_result['tag_id']
        px_per_mm = calibration_result['px_per_mm']

        # 绘制四个角点
        for i, corner in enumerate(corners):
            x, y = int(corner[0]), int(corner[1])
            cv2.circle(vis_image, (x, y), 5, (0, 255, 0), -1)
            cv2.putText(vis_image, str(i), (x+10, y),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

        # 绘制边框
        pts = corners.astype(np.int32).reshape((-1, 1, 2))
        cv2.polylines(vis_image, [pts], True, (0, 255, 0), 2)

        # 绘制中心点
        cx, cy = int(center[0]), int(center[1])
        cv2.circle(vis_image, (cx, cy), 8, (0, 0, 255), -1)

        # 绘制信息文本
        text_lines = [
            f"Tag ID: {tag_id}",
            f"Size: {calibration_result['tag_size_px']:.2f} px",
            f"Ratio: {px_per_mm:.3f} px/mm"
        ]

        y_offset = 30
        for i, line in enumerate(text_lines):
            cv2.putText(vis_image, line, (10, y_offset + i*30),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

        return vis_image

    @staticmethod
    def _compute_tag_area(corners: np.ndarray) -> float:
        """计算Tag的面积（用于多Tag选择）"""
        # Shoelace公式计算多边形面积
        x = corners[:, 0]
        y = corners[:, 1]
        return 0.5 * np.abs(np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1)))


def test_apriltag_detection(image_path: str, tag_size_mm: float = 3.76):
    """
    测试AprilTag检测功能

    参数:
        image_path: 测试图像路径
        tag_size_mm: AprilTag实际尺寸
    """
    logging.basicConfig(level=logging.INFO)

    # 读取图像
    image = cv2.imread(image_path)
    if image is None:
        logging.error(f"无法读取图像：{image_path}")
        return

    # 创建标定器
    calibrator = AprilTagCalibrator(tag_size_mm=tag_size_mm)

    # 检测AprilTag
    result = calibrator.detect(image)

    if result:
        print("\n=== AprilTag检测结果 ===")
        print(f"Tag ID: {result['tag_id']}")
        print(f"Tag尺寸（像素）: {result['tag_size_px']:.2f} px")
        print(f"像素/毫米比例: {result['px_per_mm']:.3f} px/mm")
        print(f"中心坐标: ({result['center'][0]:.2f}, {result['center'][1]:.2f})")
        print(f"角点坐标:\n{result['corners']}")

        # 可视化
        vis_image = calibrator.visualize(image, result)

        # 保存结果
        output_path = image_path.replace('.', '_apriltag_detected.')
        cv2.imwrite(output_path, vis_image)
        print(f"\n可视化结果已保存到：{output_path}")
    else:
        print("未检测到AprilTag")


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1:
        test_apriltag_detection(sys.argv[1])
    else:
        print("用法: python apriltag_calibration.py <image_path>")
