"""
开裂比分析模块

基于原shiitake_mushroom项目的开裂比计算算法
使用HSV双通道Otsu阈值法提取花纹
"""

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


class CrackRatioAnalyzer:
    """开裂比分析器（复用原项目算法）"""

    def __init__(
        self,
        enable_color_constancy: bool = True,
        enable_quantile_protection: bool = True,
        min_component_ratio: float = 0.001
    ):
        """
        初始化开裂比分析器

        参数:
            enable_color_constancy: 启用Gray-World白平衡预处理
            enable_quantile_protection: 启用Otsu阈值分位数保护
            min_component_ratio: 连通域最小面积比（相对菌盖面积），低于此值的连通域被过滤
        """
        self.enable_color_constancy = enable_color_constancy
        self.enable_quantile_protection = enable_quantile_protection
        self.min_component_ratio = min_component_ratio
        logging.info(
            f"开裂比分析器初始化完成 "
            f"(color_constancy={enable_color_constancy}, "
            f"quantile_protection={enable_quantile_protection}, "
            f"min_component_ratio={min_component_ratio})"
        )

    def _apply_gray_world(self, image: np.ndarray) -> np.ndarray:
        """
        Gray-World白平衡：将BGR三通道均值归一化到灰色世界假设

        参数:
            image: BGR格式图像

        返回:
            白平衡后的BGR图像
        """
        result = image.astype(np.float32)
        avg_b = np.mean(result[:, :, 0])
        avg_g = np.mean(result[:, :, 1])
        avg_r = np.mean(result[:, :, 2])
        avg_gray = (avg_b + avg_g + avg_r) / 3.0

        if avg_b > 0:
            result[:, :, 0] *= avg_gray / avg_b
        if avg_g > 0:
            result[:, :, 1] *= avg_gray / avg_g
        if avg_r > 0:
            result[:, :, 2] *= avg_gray / avg_r

        return np.clip(result, 0, 255).astype(np.uint8)

    def _filter_small_components(self, crack_mask: np.ndarray, cap_area_px: int) -> np.ndarray:
        """
        连通域过滤：移除面积小于菌盖面积 * min_component_ratio 的噪点连通域

        参数:
            crack_mask: 花纹二值掩码
            cap_area_px: 菌盖总面积（像素）

        返回:
            过滤后的花纹掩码
        """
        min_area = int(cap_area_px * self.min_component_ratio)
        num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(crack_mask, connectivity=8)

        filtered_mask = np.zeros_like(crack_mask)
        for i in range(1, num_labels):  # 跳过背景(label=0)
            if stats[i, cv2.CC_STAT_AREA] >= min_area:
                filtered_mask[labels == i] = 255

        return filtered_mask

    def extract_crack_pattern(self, image: np.ndarray, cap_mask: np.ndarray) -> Dict:
        """
        提取菌盖花纹并计算开裂比

        ROI优化说明：
        - 本方法已优化为ROI处理模式
        - 传入的image应该是裁剪后的ROI图像，而非全图
        - HSV转换只在ROI上进行，显著减少计算量
        - cap_mask应与ROI图像尺寸一致

        原理：
        - 花纹特征：高亮度(V) + 低饱和度(S) = 白色/奶油色
        - 底色特征：中等亮度 + 较高饱和度 = 棕色
        - 使用Otsu算法自适应确定阈值

        参数:
            image: BGR格式原始图像（推荐传入ROI而非全图）
            cap_mask: 菌盖二值掩码（0=背景，255=菌盖）

        返回:
            结果字典：
            - crack_ratio: 开裂比（0-1之间）
            - crack_ratio_percent: 开裂比百分比
            - crack_area_px: 花纹面积（像素）
            - cap_area_px: 菌盖总面积（像素）
            - crack_mask: 花纹二值掩码
            - v_threshold: V通道Otsu阈值
            - s_threshold: S通道Otsu阈值
        """
        # Gray-World白平衡预处理
        processed_image = image
        if self.enable_color_constancy:
            processed_image = self._apply_gray_world(image)

        # 转换到HSV颜色空间（只处理ROI图像，而非全图）
        hsv = cv2.cvtColor(processed_image, cv2.COLOR_BGR2HSV)
        h, s, v = cv2.split(hsv)

        # 只在菌盖区域内计算（进一步减少计算量）
        cap_region = (cap_mask > 0)

        # 提取菌盖区域的S和V通道
        s_cap = s[cap_region]
        v_cap = v[cap_region]

        # 使用Otsu算法计算自适应阈值
        v_threshold, _ = cv2.threshold(v_cap, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        s_threshold, _ = cv2.threshold(s_cap, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

        # Otsu阈值分位数保护
        if self.enable_quantile_protection:
            v_threshold = max(v_threshold, np.percentile(v_cap, 40))
            s_threshold = min(s_threshold, np.percentile(s_cap, 60))

        # 花纹条件：V > V_threshold AND S < S_threshold
        crack_mask = np.zeros_like(cap_mask)
        crack_mask[(v > v_threshold) & (s < s_threshold) & cap_region] = 255

        # 形态学处理去噪
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
        crack_mask = cv2.morphologyEx(crack_mask, cv2.MORPH_OPEN, kernel)  # 开运算去小噪点
        crack_mask = cv2.morphologyEx(crack_mask, cv2.MORPH_CLOSE, kernel)  # 闭运算填补空洞

        # 计算面积
        cap_area_px = np.sum(cap_mask > 0)

        # 连通域过滤
        if self.min_component_ratio > 0:
            crack_mask = self._filter_small_components(crack_mask, cap_area_px)

        crack_area_px = np.sum(crack_mask > 0)

        # 计算开裂比
        if cap_area_px > 0:
            crack_ratio = crack_area_px / cap_area_px
        else:
            crack_ratio = 0.0

        crack_ratio_percent = crack_ratio * 100

        result = {
            'crack_ratio': crack_ratio,
            'crack_ratio_percent': crack_ratio_percent,
            'crack_area_px': crack_area_px,
            'cap_area_px': cap_area_px,
            'crack_mask': crack_mask,
            'v_threshold': int(v_threshold),
            's_threshold': int(s_threshold)
        }

        logging.info(
            f"开裂比分析完成 - {crack_ratio_percent:.2f}% "
            f"(V_th={v_threshold}, S_th={s_threshold})"
        )

        return result

    def visualize(self, image: np.ndarray, cap_mask: np.ndarray,
                  result: Dict, alpha: float = 0.5) -> np.ndarray:
        """
        可视化开裂比分析结果

        参数:
            image: 原始BGR图像
            cap_mask: 菌盖掩码
            result: extract_crack_pattern()返回的结果
            alpha: 掩码透明度

        返回:
            可视化图像（左：原图+菌盖轮廓，右：花纹掩码）
        """
        # 创建左右两栏的可视化
        h, w = image.shape[:2]
        vis_image = np.zeros((h, w * 2, 3), dtype=np.uint8)

        # 左侧：原图 + 菌盖轮廓
        left_image = image.copy()

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

        vis_image[:, :w] = left_image

        # 右侧：花纹掩码可视化
        right_image = image.copy()

        # 创建彩色掩码覆盖层
        overlay = np.zeros_like(right_image)

        # 菌盖区域：绿色
        overlay[cap_mask > 0] = [0, 255, 0]

        # 花纹区域：红色
        overlay[result['crack_mask'] > 0] = [0, 0, 255]

        # 混合
        right_image = cv2.addWeighted(right_image, 1 - alpha, overlay, alpha, 0)

        vis_image[:, w:] = right_image

        # 绘制文本信息
        text_lines = [
            f"Crack Ratio: {result['crack_ratio_percent']:.2f}%",
            f"V_threshold: {result['v_threshold']}",
            f"S_threshold: {result['s_threshold']}"
        ]

        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, (255, 255, 255), 2)

        return vis_image

    def visualize_simple(self, image: np.ndarray, cap_mask: np.ndarray,
                        result: Dict) -> np.ndarray:
        """
        简单可视化（单图，叠加显示）

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

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

        # 创建彩色掩码
        overlay = np.zeros_like(vis_image)

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

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

        # 添加轮廓
        vis_image = cv2.add(vis_image, overlay)

        # 添加文本
        text = f"Crack: {result['crack_ratio_percent']:.1f}%"
        cv2.putText(vis_image, text, (10, 30),
                   cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 255), 2)

        return vis_image


def test_crack_analysis(image_path: str, mask_path: str):
    """
    测试开裂比分析功能

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

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

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

    # 创建分析器
    analyzer = CrackRatioAnalyzer()

    # 分析开裂比
    result = analyzer.extract_crack_pattern(image, mask)

    print("\n=== 开裂比分析结果 ===")
    print(f"开裂比: {result['crack_ratio_percent']:.2f}%")
    print(f"花纹面积: {result['crack_area_px']} px^2")
    print(f"菌盖面积: {result['cap_area_px']} px^2")
    print(f"V阈值: {result['v_threshold']}")
    print(f"S阈值: {result['s_threshold']}")

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

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


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 2:
        test_crack_analysis(sys.argv[1], sys.argv[2])
    else:
        print("用法: python crack_analyzer.py <image_path> <mask_path>")
