"""
花纹类型识别模块

根据《干香菇品质评价规范》(T/HBSYJ 001-2023)标准
实现花纹类型的自动识别和分类
"""

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


class MushroomPatternClassifier:
    """香菇花纹类型分类器"""

    def __init__(self, flatness_threshold: float = 0.65, min_circularity: float = 0.70):
        """
        初始化花纹分类器

        参数:
            flatness_threshold: 平坦度阈值（保留用于参考）
            min_circularity: 圆形度阈值，低于该值判定为残次菇
        """
        # 颜色阈值（HSV色相范围）
        # 白色花纹：低饱和度 + 高亮度
        # 茶褐色花纹：H在10-30（橙褐色范围）
        self.white_pattern_s_max = 80  # 白色花纹最大饱和度
        self.white_pattern_v_min = 100  # 白色花纹最小亮度

        self.brown_pattern_h_min = 5   # 茶褐色H最小值
        self.brown_pattern_h_max = 30  # 茶褐色H最大值

        self.flatness_threshold = flatness_threshold
        self.min_circularity = min_circularity

        logging.info("花纹分类器初始化完成")

    def _estimate_circularity(self, cap_mask: np.ndarray) -> float:
        """
        估计菌盖轮廓圆形度

        圆形度 = 4πA / P²，越接近1越接近圆形。
        """
        contours, _ = cv2.findContours(cap_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if not contours:
            return 0.0

        largest_contour = max(contours, key=cv2.contourArea)
        area = cv2.contourArea(largest_contour)
        perimeter = cv2.arcLength(largest_contour, True)
        if perimeter <= 0:
            return 0.0

        return 4 * np.pi * area / (perimeter * perimeter)

    def _estimate_flatness(
        self,
        image: np.ndarray,
        cap_mask: np.ndarray
    ) -> float:
        """
        估计菌盖平坦度（用于板菇判定）

        通过轮廓凸包面积比和边缘亮度梯度标准差综合评估。
        板菇特征：高凸包率（形状规整）+ 低梯度std（反光均匀）。

        参数:
            image: BGR格式ROI图像
            cap_mask: 菌盖二值掩码

        返回:
            flatness_score: 0-1之间，越高越可能是板菇
        """
        # 1. 轮廓凸包面积比
        contours, _ = cv2.findContours(cap_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if not contours:
            return 0.0

        largest_contour = max(contours, key=cv2.contourArea)
        contour_area = cv2.contourArea(largest_contour)
        hull = cv2.convexHull(largest_contour)
        hull_area = cv2.contourArea(hull)

        convexity = contour_area / hull_area if hull_area > 0 else 0.0

        # 2. 边缘亮度梯度标准差
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        # 在菌盖区域内计算Sobel梯度
        grad_x = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
        grad_y = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
        grad_mag = np.sqrt(grad_x**2 + grad_y**2)

        cap_region = cap_mask > 0
        if np.sum(cap_region) == 0:
            return 0.0

        grad_in_cap = grad_mag[cap_region]
        grad_std = np.std(grad_in_cap)

        # 归一化梯度std到0-1（std越低越平坦，分数越高）
        # 经验值：grad_std < 15 非常平坦，> 60 纹理丰富
        grad_score = max(0.0, min(1.0, 1.0 - (grad_std - 15) / 45.0))

        # 综合评分：凸包率权重0.4 + 梯度平坦度权重0.6
        flatness_score = 0.4 * convexity + 0.6 * grad_score

        return flatness_score

    def classify_pattern(
        self,
        image: np.ndarray,
        cap_mask: np.ndarray,
        crack_mask: np.ndarray,
        crack_ratio: float
    ) -> Dict:
        """
        识别花纹类型

        参数:
            image: BGR格式原始图像（ROI）
            cap_mask: 菌盖二值掩码
            crack_mask: 花纹二值掩码
            crack_ratio: 花纹比例（0-1）

        返回:
            分类结果字典：
            - pattern_type: 花纹类型（天白花菇/白花菇3A/2A/A/茶花菇/板菇/残次菇）
            - pattern_category: 花纹大类（白花纹/茶花纹/无花纹/残次）
            - white_pattern_ratio: 白色花纹占比
            - brown_pattern_ratio: 茶褐色花纹占比
            - crack_ratio_percent: 总花纹占比百分比
            - flatness_score: 平坦度评分（仅低花纹时计算）
            - suspected_ban_gu: 平坦度是否超过阈值（保留字段）
        """
        result = {
            'crack_ratio_percent': crack_ratio * 100
        }

        # 先做残次菇过滤（轮廓不近似圆形）
        circularity = self._estimate_circularity(cap_mask)
        result['circularity'] = circularity
        if circularity < self.min_circularity:
            result['white_pattern_ratio'] = 0.0
            result['brown_pattern_ratio'] = 0.0
            result['flatness_score'] = 0.0
            result['suspected_ban_gu'] = False
            result['pattern_type'] = '残次菇'
            result['pattern_category'] = '残次'
            logging.info(
                f"识别为残次菇（圆形度: {circularity:.3f} < {self.min_circularity:.2f}）"
            )
            return result

        # 如果花纹比例很低（<5%），统一判定为板菇
        if crack_ratio < 0.05:
            flatness_score = self._estimate_flatness(image, cap_mask)
            suspected_ban_gu = flatness_score > self.flatness_threshold

            result['white_pattern_ratio'] = 0.0
            result['brown_pattern_ratio'] = 0.0
            result['flatness_score'] = flatness_score
            result['suspected_ban_gu'] = suspected_ban_gu

            result['pattern_type'] = '板菇'
            result['pattern_category'] = '无花纹'
            logging.info(
                f"识别为板菇（花纹比例: {crack_ratio*100:.2f}%, "
                f"平坦度: {flatness_score:.3f}）"
            )
            return result

        # 转换到HSV颜色空间
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        h, s, v = cv2.split(hsv)

        # 只分析花纹区域
        crack_region = (crack_mask > 0)

        if np.sum(crack_region) == 0:
            result['pattern_type'] = '板菇'
            result['pattern_category'] = '无花纹'
            result['white_pattern_ratio'] = 0.0
            result['brown_pattern_ratio'] = 0.0
            return result

        # 提取花纹区域的HSV值
        h_crack = h[crack_region]
        s_crack = s[crack_region]
        v_crack = v[crack_region]

        # 识别白色花纹：低饱和度 + 高亮度
        white_pattern_mask = (
            (s_crack < self.white_pattern_s_max) &
            (v_crack > self.white_pattern_v_min)
        )
        white_pattern_count = np.sum(white_pattern_mask)

        # 识别茶褐色花纹：H在橙褐色范围
        brown_pattern_mask = (
            (h_crack >= self.brown_pattern_h_min) &
            (h_crack <= self.brown_pattern_h_max)
        )
        brown_pattern_count = np.sum(brown_pattern_mask)

        # 计算占菌盖总面积的比例
        cap_area = np.sum(cap_mask > 0)
        white_pattern_ratio = white_pattern_count / cap_area if cap_area > 0 else 0
        brown_pattern_ratio = brown_pattern_count / cap_area if cap_area > 0 else 0

        result['white_pattern_ratio'] = white_pattern_ratio
        result['brown_pattern_ratio'] = brown_pattern_ratio

        # 根据标准进行分类
        pattern_type, pattern_category = self._classify_by_standard(
            white_pattern_ratio, brown_pattern_ratio, crack_ratio
        )

        result['pattern_type'] = pattern_type
        result['pattern_category'] = pattern_category
        result['flatness_score'] = 0.0
        result['suspected_ban_gu'] = False

        logging.info(
            f"花纹识别完成 - 类型: {pattern_type}, "
            f"白色花纹: {white_pattern_ratio*100:.2f}%, "
            f"茶褐色花纹: {brown_pattern_ratio*100:.2f}%"
        )

        return result

    def _classify_by_standard(
        self,
        white_ratio: float,
        brown_ratio: float,
        total_crack_ratio: float
    ) -> Tuple[str, str]:
        """
        根据《干香菇品质评价规范》标准进行分类

        参数:
            white_ratio: 白色花纹占菌盖面积比例
            brown_ratio: 茶褐色花纹占菌盖面积比例
            total_crack_ratio: 总花纹比例

        返回:
            (花纹类型, 花纹大类)
        """
        white_percent = white_ratio * 100
        brown_percent = brown_ratio * 100

        # 判断主要花纹类型
        if white_ratio > brown_ratio:
            # 白色花纹为主
            if white_percent >= 80:
                return '天白花菇', '白花纹'
            elif white_percent >= 60:
                return '白花菇3A', '白花纹'
            elif white_percent >= 26:
                return '白花菇2A', '白花纹'
            elif white_percent >= 5:
                return '白花菇A', '白花纹'
            else:
                return '板菇', '无花纹'
        else:
            # 茶褐色花纹为主
            if brown_percent >= 5:
                return '茶花菇', '茶花纹'
            else:
                # 有花纹但颜色不明显，归为板菇（与光面菇合并）
                return '板菇', '无花纹'

    def visualize(
        self,
        image: np.ndarray,
        cap_mask: np.ndarray,
        crack_mask: np.ndarray,
        result: Dict
    ) -> np.ndarray:
        """
        可视化花纹分类结果

        参数:
            image: 原始BGR图像
            cap_mask: 菌盖掩码
            crack_mask: 花纹掩码
            result: classify_pattern()返回的结果

        返回:
            可视化图像
        """
        vis_image = image.copy()

        # 绘制菌盖轮廓（绿色）
        contours, _ = cv2.findContours(cap_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        cv2.drawContours(vis_image, contours, -1, (0, 255, 0), 2)

        # 绘制花纹区域（半透明红色）
        crack_overlay = np.zeros_like(vis_image)
        crack_overlay[crack_mask > 0] = [0, 0, 255]
        vis_image = cv2.addWeighted(vis_image, 0.7, crack_overlay, 0.3, 0)

        # 添加分类信息文本
        pattern_type = result.get('pattern_type', '未知')
        white_ratio = result.get('white_pattern_ratio', 0) * 100
        brown_ratio = result.get('brown_pattern_ratio', 0) * 100

        # 在图像上方添加文本
        y_offset = 30
        cv2.putText(vis_image, f"Type: {pattern_type}", (10, y_offset),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)

        y_offset += 35
        cv2.putText(vis_image, f"White: {white_ratio:.1f}%", (10, y_offset),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)

        y_offset += 35
        cv2.putText(vis_image, f"Brown: {brown_ratio:.1f}%", (10, y_offset),
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)

        return vis_image


def test_pattern_classifier(image_path: str, mask_path: str, crack_mask_path: str):
    """
    测试花纹分类功能

    参数:
        image_path: 原始图像路径
        mask_path: 菌盖掩码路径
        crack_mask_path: 花纹掩码路径
    """
    logging.basicConfig(level=logging.INFO)

    # 读取图像和掩码
    image = cv2.imread(image_path)
    cap_mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
    crack_mask = cv2.imread(crack_mask_path, cv2.IMREAD_GRAYSCALE)

    if image is None or cap_mask is None or crack_mask is None:
        logging.error("无法读取图像或掩码")
        return

    # 计算花纹比例
    cap_area = np.sum(cap_mask > 0)
    crack_area = np.sum(crack_mask > 0)
    crack_ratio = crack_area / cap_area if cap_area > 0 else 0

    # 创建分类器
    classifier = MushroomPatternClassifier()

    # 分类
    result = classifier.classify_pattern(image, cap_mask, crack_mask, crack_ratio)

    print("\n=== 花纹分类结果 ===")
    print(f"花纹类型: {result['pattern_type']}")
    print(f"花纹大类: {result['pattern_category']}")
    print(f"白色花纹占比: {result['white_pattern_ratio']*100:.2f}%")
    print(f"茶褐色花纹占比: {result['brown_pattern_ratio']*100:.2f}%")
    print(f"总花纹占比: {result['crack_ratio_percent']:.2f}%")

    # 可视化
    vis_image = classifier.visualize(image, cap_mask, crack_mask, result)

    # 保存结果
    output_path = image_path.replace('.', '_pattern_classification.')
    cv2.imwrite(output_path, vis_image)
    print(f"\n可视化结果已保存到：{output_path}")


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 3:
        test_pattern_classifier(sys.argv[1], sys.argv[2], sys.argv[3])
    else:
        print("用法: python pattern_classifier.py <image_path> <cap_mask_path> <crack_mask_path>")
