"""
菌盖最大直径计算模块

功能：
1. 基于分割掩码提取轮廓
2. 计算凸包
3. 使用旋转卡壳算法计算最大直径
4. 提供可视化和像素到毫米的转换
"""

import cv2
import numpy as np
from scipy.spatial import ConvexHull
from typing import Tuple, Optional, Dict
import logging


class DiameterCalculator:
    """菌盖直径计算器"""

    def __init__(self, method: str = 'rotating_calipers'):
        """
        初始化直径计算器

        参数:
            method: 计算方法
                - 'rotating_calipers': 旋转卡壳算法（最精确，计算真实最远点对）
                - 'min_rect': 最小外接矩形长轴（较快，但可能略小）
        """
        if method not in ['rotating_calipers', 'min_rect']:
            raise ValueError(f"不支持的方法: {method}")

        self.method = method
        logging.info(f"直径计算器初始化完成，方法={method}")

    def calculate(self, mask: np.ndarray, px_per_mm: Optional[float] = None,
                  homography: Optional[np.ndarray] = None,
                  offset: Optional[Tuple[int, int]] = None) -> Dict:
        """
        计算菌盖最大直径（支持透视校正）

        参数:
            mask: 二值掩码图像（0=背景，255=菌盖）
            px_per_mm: 像素/毫米比例（可选，用于转换为毫米）
            homography: 单应性矩阵（可选，用于透视校正）
            offset: ROI在全图中的偏移量 (offset_x, offset_y)，用于坐标映射

        返回:
            结果字典：
            - diameter_px: 直径（像素，未校正）
            - diameter_px_corrected: 直径（像素，透视校正后）
            - diameter_mm: 直径（毫米），如果提供了px_per_mm
            - point1: 直径端点1坐标 (x, y)
            - point2: 直径端点2坐标 (x, y)
            - center: 菌盖中心坐标
            - cap_area_px: 菌盖面积（像素）
            - cap_area_mm2: 菌盖面积（平方毫米），如果提供了px_per_mm
            - contour: 外轮廓点集
        """
        # 提取轮廓
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        if len(contours) == 0:
            logging.error("未找到轮廓")
            return None

        # 选择最大轮廓
        contour = max(contours, key=cv2.contourArea)
        points = contour.squeeze()

        # 处理特殊情况
        if len(points.shape) == 1:
            points = points.reshape(-1, 2)

        # 计算面积
        cap_area_px = cv2.contourArea(contour)

        # 计算中心（矩的方法）
        M = cv2.moments(contour)
        if M['m00'] != 0:
            cx = M['m10'] / M['m00']
            cy = M['m01'] / M['m00']
            center = (cx, cy)
        else:
            center = tuple(points.mean(axis=0))

        # 根据方法计算直径
        if self.method == 'rotating_calipers':
            diameter_px, pt1, pt2 = self._rotating_calipers(points)
        else:  # min_rect
            diameter_px, pt1, pt2 = self._min_area_rect(contour)

        # 构建结果
        result = {
            'diameter_px': diameter_px,
            'point1': pt1,
            'point2': pt2,
            'center': center,
            'cap_area_px': cap_area_px,
            'contour': contour,
            'method': self.method
        }

        # ========== 透视校正（关键优化） ==========
        if homography is not None:
            # 如果是ROI坐标，需要先映射回全图坐标
            if offset is not None:
                offset_x, offset_y = offset
                pt1_global = (pt1[0] + offset_x, pt1[1] + offset_y)
                pt2_global = (pt2[0] + offset_x, pt2[1] + offset_y)
            else:
                pt1_global = pt1
                pt2_global = pt2

            corrected_result = self._correct_with_homography(
                diameter_px, pt1_global, pt2_global, cap_area_px, homography
            )
            result.update(corrected_result)

            # 使用校正后的直径计算毫米值
            if px_per_mm is not None:
                result['diameter_mm'] = result['diameter_px_corrected'] / px_per_mm
                result['cap_area_mm2'] = result['cap_area_px_corrected'] / (px_per_mm ** 2)
                result['px_per_mm'] = px_per_mm

            logging.info(
                f"直径计算完成 - 原始:{diameter_px:.2f}px, "
                f"校正:{result['diameter_px_corrected']:.2f}px"
                + (f" ({result['diameter_mm']:.2f}mm)" if px_per_mm else "")
            )
        else:
            # 如果提供了像素/毫米比例，计算毫米值（未校正）
            if px_per_mm is not None:
                result['diameter_mm'] = diameter_px / px_per_mm
                result['cap_area_mm2'] = cap_area_px / (px_per_mm ** 2)
                result['px_per_mm'] = px_per_mm

            logging.info(
                f"直径计算完成 - {diameter_px:.2f}px"
                + (f" ({result['diameter_mm']:.2f}mm)" if px_per_mm else "")
            )

        return result

    def _correct_with_homography(self, diameter_px: float,
                                  pt1: Tuple[float, float],
                                  pt2: Tuple[float, float],
                                  area_px: float,
                                  H: np.ndarray) -> Dict:
        """
        使用单应性矩阵校正直径和面积

        原理：
        - 将直径端点从畸变图像变换到校正图像
        - 重新计算校正后的直径
        - 估算校正后的面积

        参数:
            diameter_px: 原始直径（像素）
            pt1, pt2: 直径端点
            area_px: 原始面积
            H: 单应性矩阵

        返回:
            包含校正值的字典
        """
        # 将端点转换为齐次坐标
        pt1_homogeneous = np.array([[pt1[0], pt1[1]]], dtype=np.float32).reshape(-1, 1, 2)
        pt2_homogeneous = np.array([[pt2[0], pt2[1]]], dtype=np.float32).reshape(-1, 1, 2)

        # 透视变换：从畸变图像 → 校正图像
        pt1_corrected = cv2.perspectiveTransform(pt1_homogeneous, H).reshape(2)
        pt2_corrected = cv2.perspectiveTransform(pt2_homogeneous, H).reshape(2)

        # 计算校正后的直径
        diameter_px_corrected = np.linalg.norm(pt1_corrected - pt2_corrected)

        # 估算校正后的面积
        # 使用单应性矩阵的行列式作为面积缩放因子
        scale_factor = np.abs(np.linalg.det(H[:2, :2]))
        area_px_corrected = area_px * scale_factor

        return {
            'diameter_px_corrected': diameter_px_corrected,
            'cap_area_px_corrected': area_px_corrected,
            'point1_corrected': tuple(pt1_corrected),
            'point2_corrected': tuple(pt2_corrected),
            'correction_ratio': diameter_px_corrected / diameter_px
        }

    def _rotating_calipers(self, points: np.ndarray) -> Tuple[float, Tuple, Tuple]:
        """
        旋转卡壳算法计算凸包最远点对

        参数:
            points: 轮廓点集 (N, 2)

        返回:
            (最大直径, 端点1, 端点2)
        """
        try:
            # 计算凸包
            hull = ConvexHull(points)
            hull_points = points[hull.vertices]

            # 暴力搜索凸包上的最远点对（旋转卡壳的简化版）
            max_dist = 0
            pt1_best = None
            pt2_best = None

            n = len(hull_points)
            for i in range(n):
                for j in range(i + 1, n):
                    dist = np.linalg.norm(hull_points[i] - hull_points[j])
                    if dist > max_dist:
                        max_dist = dist
                        pt1_best = tuple(hull_points[i].astype(float))
                        pt2_best = tuple(hull_points[j].astype(float))

            return max_dist, pt1_best, pt2_best

        except Exception as e:
            logging.warning(f"旋转卡壳算法失败: {e}，回退到最小外接矩形方法")
            # 回退方案
            contour = points.reshape(-1, 1, 2).astype(np.int32)
            return self._min_area_rect(contour)

    def _min_area_rect(self, contour: np.ndarray) -> Tuple[float, Tuple, Tuple]:
        """
        使用最小外接矩形的长轴作为直径

        参数:
            contour: OpenCV格式轮廓

        返回:
            (最大直径, 端点1, 端点2)
        """
        rect = cv2.minAreaRect(contour)
        box = cv2.boxPoints(rect)

        # 找到矩形的最长边
        side1 = np.linalg.norm(box[0] - box[1])
        side2 = np.linalg.norm(box[1] - box[2])

        if side1 > side2:
            diameter = side1
            pt1 = tuple(((box[0] + box[1]) / 2).astype(float))
            pt2 = tuple(((box[2] + box[3]) / 2).astype(float))
        else:
            diameter = side2
            pt1 = tuple(((box[1] + box[2]) / 2).astype(float))
            pt2 = tuple(((box[3] + box[0]) / 2).astype(float))

        return diameter, pt1, pt2

    def visualize(self, image: np.ndarray, result: Dict,
                  show_contour: bool = True,
                  show_diameter_line: bool = True) -> np.ndarray:
        """
        可视化直径计算结果

        参数:
            image: 原始BGR图像
            result: calculate()返回的结果字典
            show_contour: 是否显示轮廓
            show_diameter_line: 是否显示直径线段

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

        if result is None:
            return vis_image

        # 绘制轮廓
        if show_contour and 'contour' in result:
            cv2.drawContours(vis_image, [result['contour']], -1, (0, 255, 0), 2)

        # 绘制直径线段
        if show_diameter_line:
            pt1 = tuple(map(int, result['point1']))
            pt2 = tuple(map(int, result['point2']))

            # 绘制线段
            cv2.line(vis_image, pt1, pt2, (0, 0, 255), 3)

            # 绘制端点
            cv2.circle(vis_image, pt1, 6, (255, 0, 0), -1)
            cv2.circle(vis_image, pt2, 6, (255, 0, 0), -1)

        # 绘制中心点
        center = tuple(map(int, result['center']))
        cv2.circle(vis_image, center, 5, (255, 255, 0), -1)

        # 绘制文本信息
        text_lines = []

        if 'diameter_mm' in result:
            text_lines.append(f"Diameter: {result['diameter_mm']:.2f} mm")
        else:
            text_lines.append(f"Diameter: {result['diameter_px']:.2f} px")

        if 'cap_area_mm2' in result:
            text_lines.append(f"Area: {result['cap_area_mm2']:.2f} mm^2")

        # 在图像上方绘制文本
        y_offset = 30
        for i, line in enumerate(text_lines):
            cv2.putText(vis_image, line, (10, y_offset + i * 35),
                       cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 255), 2)

        return vis_image


def test_diameter_calculation(mask_path: str, px_per_mm: float = 26.6):
    """
    测试直径计算功能

    参数:
        mask_path: 二值掩码图像路径
        px_per_mm: 像素/毫米比例
    """
    logging.basicConfig(level=logging.INFO)

    # 读取掩码
    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
    if mask is None:
        logging.error(f"无法读取掩码：{mask_path}")
        return

    # 创建计算器
    calc = DiameterCalculator(method='rotating_calipers')

    # 计算直径
    result = calc.calculate(mask, px_per_mm=px_per_mm)

    if result:
        print("\n=== 直径计算结果 ===")
        print(f"直径（像素）: {result['diameter_px']:.2f} px")
        if 'diameter_mm' in result:
            print(f"直径（毫米）: {result['diameter_mm']:.2f} mm")
        print(f"面积（像素）: {result['cap_area_px']:.2f} px^2")
        if 'cap_area_mm2' in result:
            print(f"面积（毫米）: {result['cap_area_mm2']:.2f} mm^2")
        print(f"中心坐标: ({result['center'][0]:.2f}, {result['center'][1]:.2f})")

        # 可视化（需要原图，这里用掩码代替）
        vis_image = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
        vis_image = calc.visualize(vis_image, result)

        # 保存结果
        output_path = mask_path.replace('.', '_diameter.')
        cv2.imwrite(output_path, vis_image)
        print(f"\n可视化结果已保存到：{output_path}")
    else:
        print("直径计算失败")


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1:
        test_diameter_calculation(sys.argv[1])
    else:
        print("用法: python diameter_calculator.py <mask_path>")
