"""
集成分析器

整合AprilTag标定、香菇检测、分割、直径计算和开裂比分析
"""

import cv2
import numpy as np
from pathlib import Path
from typing import Dict, List, Optional
from PIL import Image, ImageDraw, ImageFont
import logging

try:
    from .apriltag_calibration import AprilTagCalibrator
    from .mushroom_detector import MushroomDetector, MushroomSegmenter
    from .diameter_calculator import DiameterCalculator
    from .crack_analyzer import CrackRatioAnalyzer
    from .pattern_classifier import MushroomPatternClassifier
    from .parallel_processor import process_mushrooms_parallel
except ImportError:  # pragma: no cover - fallback for direct script usage
    from apriltag_calibration import AprilTagCalibrator
    from mushroom_detector import MushroomDetector, MushroomSegmenter
    from diameter_calculator import DiameterCalculator
    from crack_analyzer import CrackRatioAnalyzer
    from pattern_classifier import MushroomPatternClassifier
    from parallel_processor import process_mushrooms_parallel


class IntegratedMushroomAnalyzer:
    """香菇完整分析系统"""

    def __init__(
        self,
        yolo_model_path: str,
        sam_model_path: str,
        apriltag_size_mm: float = 37.58,
        device: str = 'cuda',
        enable_parallel: bool = False,
        max_workers: int = 4
    ):
        """
        初始化集成分析器

        参数:
            yolo_model_path: YOLO模型路径
            sam_model_path: SAM模型路径
            apriltag_size_mm: AprilTag尺寸（毫米）
            device: 运行设备
            enable_parallel: 是否启用并行处理（多个香菇同时处理）
            max_workers: 并行线程数（仅当enable_parallel=True时生效）
        """
        # 初始化各个模块
        self.apriltag_calibrator = AprilTagCalibrator(tag_size_mm=apriltag_size_mm)
        self.mushroom_detector = MushroomDetector(yolo_model_path)
        self.mushroom_segmenter = MushroomSegmenter(sam_model_path, device=device)
        self.diameter_calculator = DiameterCalculator(method='rotating_calipers')
        self.crack_analyzer = CrackRatioAnalyzer()
        self.pattern_classifier = MushroomPatternClassifier()

        self.enable_parallel = enable_parallel
        self.max_workers = max_workers

        logging.info(f"集成分析器初始化完成（并行模式: {enable_parallel}）")

    def analyze_image(self, image_path: str, visualize: bool = True) -> Dict:
        """
        分析单张图像（采用ROI处理哲学）

        处理流程（分层并行优化）：
        1. AprilTag检测与标定（全图快速扫描）
        2. YOLO检测香菇（全图快速扫描）
        3. 裁剪ROI区域（关键优化：只处理必要区域）
        4. SAM分割菌盖（仅在ROI上运行，快10-20倍）
        5. 计算最大直径（仅在ROI掩码上计算）
        6. 计算开裂比（仅在ROI图像上做HSV转换）

        性能优化说明：
        - AprilTag区域：约100×100像素
        - 每个香菇ROI：约500×500像素
        - 全图：4000×3000 = 1200万像素
        - 实际处理：10K + 3×250K ≈ 76万像素（减少94%计算量）

        参数:
            image_path: 图像路径
            visualize: 是否生成可视化图像

        返回:
            分析结果字典
        """
        logging.info(f"开始分析图像：{image_path}")

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

        result = {
            'filename': Path(image_path).name,
            'image_path': image_path,
            'success': False
        }

        # Step 1: AprilTag标定（含透视校正）
        logging.info("Step 1: AprilTag标定")
        calibration = self.apriltag_calibrator.detect(image, use_homography=True)

        if not calibration or not calibration['detected']:
            logging.error("AprilTag检测失败，无法进行标定")
            result['error'] = 'AprilTag未检测到'
            return result

        # 优先使用校正后的px/mm比例
        px_per_mm = calibration.get('px_per_mm_corrected', calibration['px_per_mm'])
        homography = calibration.get('homography')

        result['apriltag_detected'] = True
        result['px_per_mm'] = px_per_mm
        result['px_per_mm_raw'] = calibration['px_per_mm']
        result['apriltag_id'] = calibration['tag_id']
        result['reprojection_error'] = calibration.get('reprojection_error', 0)
        result['distortion_ratio'] = calibration.get('distortion_ratio', 0)
        result['homography_enabled'] = homography is not None

        if homography is not None:
            logging.info(f"透视校正已启用（重投影误差: {result['reprojection_error']:.3f}px）")

        # Step 2: YOLO检测香菇
        logging.info("Step 2: 检测香菇")
        detections = self.mushroom_detector.detect(image)

        if len(detections) == 0:
            logging.warning("未检测到香菇")
            result['error'] = '未检测到香菇'
            return result

        result['mushroom_count'] = len(detections)

        # 处理所有检测到的香菇（支持并行模式）
        if self.enable_parallel and len(detections) > 1:
            logging.info(f"使用并行模式处理 {len(detections)} 个香菇")
            mushroom_results = process_mushrooms_parallel(
                detections=detections,
                image=image,
                detector=self.mushroom_detector,
                segmenter=self.mushroom_segmenter,
                diameter_calc=self.diameter_calculator,
                crack_analyzer=self.crack_analyzer,
                px_per_mm=px_per_mm,
                max_workers=self.max_workers
            )
        else:
            # 串行处理（默认模式）
            mushroom_results = []

            for idx, detection in enumerate(detections):
                logging.info(f"处理香菇 {idx + 1}/{len(detections)}")

                mushroom_result = {
                    'mushroom_id': idx,
                    'detection_confidence': detection['confidence']
                }

                box = detection['box']

                # ========== ROI处理哲学：只处理必要区域 ==========

                # Step 2.5: 裁剪ROI（关键优化）
                logging.info("Step 2.5: 裁剪ROI区域")
                roi_image, offset = self.mushroom_detector.crop_roi(image, box, padding_ratio=0.1)

                logging.info(f"  原图尺寸: {image.shape[:2]}, ROI尺寸: {roi_image.shape[:2]} "
                            f"(减少 {100 * (1 - roi_image.size / image.size):.1f}% 计算量)")

                # Step 3: SAM分割菌盖（只在ROI上运行，快10-20倍！）
                logging.info("Step 3: 分割菌盖（ROI模式）")
                cap_mask_roi = self.mushroom_segmenter.segment_roi(roi_image, box, offset)

                if cap_mask_roi is None:
                    logging.warning(f"香菇 {idx} 分割失败")
                    mushroom_result['error'] = '分割失败'
                    mushroom_results.append(mushroom_result)
                    continue

                # Step 4: 计算最大直径（支持透视校正）
                logging.info("Step 4: 计算最大直径（透视校正模式）")
                diameter_result = self.diameter_calculator.calculate(
                    cap_mask_roi,
                    px_per_mm=px_per_mm,
                    homography=homography,  # 传入单应性矩阵进行校正
                    offset=offset  # 传入ROI偏移量，用于坐标映射
                )

                if diameter_result is None:
                    logging.warning(f"香菇 {idx} 直径计算失败")
                    mushroom_result['error'] = '直径计算失败'
                    mushroom_results.append(mushroom_result)
                    continue

                # 保存结果（优先使用校正后的值）
                mushroom_result['diameter_px'] = diameter_result.get('diameter_px_corrected',
                                                                      diameter_result['diameter_px'])
                mushroom_result['diameter_mm'] = diameter_result['diameter_mm']
                mushroom_result['cap_area_px'] = diameter_result.get('cap_area_px_corrected',
                                                                     diameter_result['cap_area_px'])
                mushroom_result['cap_area_mm2'] = diameter_result['cap_area_mm2']

                # 保存校正信息
                if 'correction_ratio' in diameter_result:
                    mushroom_result['correction_ratio'] = diameter_result['correction_ratio']

                # Step 5: 计算开裂比（只在ROI图像上做HSV转换）
                logging.info("Step 5: 计算开裂比（ROI模式）")
                crack_result = self.crack_analyzer.extract_crack_pattern(roi_image, cap_mask_roi)

                mushroom_result['crack_ratio'] = crack_result['crack_ratio']
                mushroom_result['crack_ratio_percent'] = crack_result['crack_ratio_percent']
                mushroom_result['crack_area_px'] = crack_result['crack_area_px']
                mushroom_result['v_threshold'] = crack_result['v_threshold']
                mushroom_result['s_threshold'] = crack_result['s_threshold']

                # Step 6: 花纹类型识别（符合T/HBSYJ 001-2023标准）
                logging.info("Step 6: 花纹类型识别")
                pattern_result = self.pattern_classifier.classify_pattern(
                    roi_image,
                    cap_mask_roi,
                    crack_result['crack_mask'],
                    crack_result['crack_ratio']
                )

                mushroom_result['pattern_type'] = pattern_result['pattern_type']
                mushroom_result['pattern_category'] = pattern_result['pattern_category']
                mushroom_result['white_pattern_ratio'] = pattern_result['white_pattern_ratio']
                mushroom_result['brown_pattern_ratio'] = pattern_result['brown_pattern_ratio']
                mushroom_result['flatness_score'] = pattern_result.get('flatness_score', 0.0)
                mushroom_result['suspected_ban_gu'] = pattern_result.get('suspected_ban_gu', False)

                # 保存中间结果用于可视化
                if visualize:
                    mushroom_result['_box'] = box
                    mushroom_result['_offset'] = offset  # ROI偏移量
                    mushroom_result['_cap_mask_roi'] = cap_mask_roi  # ROI掩码
                    mushroom_result['_crack_mask_roi'] = crack_result['crack_mask']  # ROI花纹掩码

                    # 直径端点需要映射回全图坐标
                    pt1_roi = diameter_result['point1']
                    pt2_roi = diameter_result['point2']
                    mushroom_result['_diameter_pt1'] = (pt1_roi[0] + offset[0], pt1_roi[1] + offset[1])
                    mushroom_result['_diameter_pt2'] = (pt2_roi[0] + offset[0], pt2_roi[1] + offset[1])

                mushroom_results.append(mushroom_result)

                logging.info(
                    f"香菇 {idx} 分析完成 - "
                    f"直径: {mushroom_result['diameter_mm']:.2f}mm, "
                    f"开裂比: {mushroom_result['crack_ratio_percent']:.2f}%, "
                    f"类型: {mushroom_result['pattern_type']}"
                )

        result['mushrooms'] = mushroom_results
        result['success'] = len(mushroom_results) > 0

        # 可视化
        if visualize and result['success']:
            result['visualization'] = self._create_visualization(
                image, calibration, mushroom_results
            )

        return result

    def _create_visualization(
        self,
        image: np.ndarray,
        calibration: Dict,
        mushroom_results: List[Dict]
    ) -> np.ndarray:
        """
        创建聚焦式可视化图像（只显示香菇ROI区域）
        使用PIL绘制中文文字，避免乱码
        """
        # 加载简体中文字体（明确指定SC - Simplified Chinese）
        try:
            font_path = '/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc'
            # TTC文件中字体顺序：0=JP, 1=HK, 2=KR, 3=SC, 4=TC
            # index=3 是简体中文 (SC)
            pil_font_large = ImageFont.truetype(font_path, 36, index=3)
            pil_font_small = ImageFont.truetype(font_path, 28, index=3)
        except:
            pil_font_large = ImageFont.load_default()
            pil_font_small = ImageFont.load_default()

        roi_visualizations = []

        for mushroom in mushroom_results:
            if 'error' in mushroom:
                continue

            offset = mushroom['_offset']
            cap_mask_roi = mushroom['_cap_mask_roi']
            crack_mask_roi = mushroom['_crack_mask_roi']

            # 用offset裁剪原图，与掩码坐标系一致
            roi_h, roi_w = cap_mask_roi.shape[:2]
            ox, oy = offset
            roi_image = image[oy:oy+roi_h, ox:ox+roi_w].copy()

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

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

            # 绘制直径线（全图坐标转ROI坐标）
            pt1_global = mushroom['_diameter_pt1']
            pt2_global = mushroom['_diameter_pt2']
            pt1_local = (int(pt1_global[0] - ox), int(pt1_global[1] - oy))
            pt2_local = (int(pt2_global[0] - ox), int(pt2_global[1] - oy))

            cv2.line(roi_image, pt1_local, pt2_local, (0, 255, 255), 4)
            cv2.circle(roi_image, pt1_local, 8, (0, 0, 255), -1)
            cv2.circle(roi_image, pt2_local, 8, (0, 0, 255), -1)

            # 用PIL绘制中文标注
            pil_image = Image.fromarray(cv2.cvtColor(roi_image, cv2.COLOR_BGR2RGB))
            draw = ImageDraw.Draw(pil_image, 'RGBA')

            # 半透明黑色背景框
            box_h = 110
            draw.rectangle([(8, 8), (roi_w - 8, 8 + box_h)],
                          fill=(0, 0, 0, 180), outline=(255, 255, 255, 200), width=2)

            # 第一行：直径（大字）
            line1 = f"直径: {mushroom['diameter_mm']:.2f} mm"
            draw.text((20, 16), line1, font=pil_font_large, fill=(255, 255, 100))

            # 第二行：花纹类型 + 花纹比例
            line2 = f"{mushroom['pattern_type']}  花纹: {mushroom['crack_ratio_percent']:.1f}%"
            draw.text((20, 64), line2, font=pil_font_small, fill=(100, 255, 100))

            roi_image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
            roi_visualizations.append(roi_image)

        if len(roi_visualizations) == 0:
            return image

        if len(roi_visualizations) == 1:
            return roi_visualizations[0]

        # 多个香菇：统一高度后水平拼接
        max_height = max(roi.shape[0] for roi in roi_visualizations)
        resized_rois = []
        for roi in roi_visualizations:
            if roi.shape[0] != max_height:
                scale = max_height / roi.shape[0]
                new_w = int(roi.shape[1] * scale)
                roi = cv2.resize(roi, (new_w, max_height))
            resized_rois.append(roi)

        return np.hstack(resized_rois)


def test_integrated_analyzer(
    image_path: str,
    yolo_model: str,
    sam_model: str,
    output_dir: str = './outputs'
):
    """
    测试集成分析器

    参数:
        image_path: 测试图像路径
        yolo_model: YOLO模型路径
        sam_model: SAM模型路径
        output_dir: 输出目录
    """
    logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')

    # 创建分析器
    analyzer = IntegratedMushroomAnalyzer(
        yolo_model_path=yolo_model,
        sam_model_path=sam_model,
        apriltag_size_mm=37.58
    )

    # 分析图像
    result = analyzer.analyze_image(image_path, visualize=True)

    if result and result['success']:
        print("\n" + "="*60)
        print("分析结果")
        print("="*60)
        print(f"文件名: {result['filename']}")
        print(f"AprilTag检测: {'✓' if result['apriltag_detected'] else '✗'}")
        print(f"像素/毫米比例: {result['px_per_mm']:.3f} px/mm")
        print(f"检测到香菇数量: {result['mushroom_count']}")
        print()

        for mushroom in result['mushrooms']:
            if 'error' not in mushroom:
                print(f"香菇 #{mushroom['mushroom_id']}")
                print(f"  直径: {mushroom['diameter_mm']:.2f} mm")
                print(f"  面积: {mushroom['cap_area_mm2']:.2f} mm²")
                print(f"  开裂比: {mushroom['crack_ratio_percent']:.2f}%")
                print(f"  花纹类型: {mushroom['pattern_type']}")
                print(f"  白色花纹: {mushroom['white_pattern_ratio']*100:.2f}%")
                print(f"  茶褐色花纹: {mushroom['brown_pattern_ratio']*100:.2f}%")
                print(f"  置信度: {mushroom['detection_confidence']:.3f}")
                print()

        # 保存可视化
        if 'visualization' in result:
            Path(output_dir).mkdir(parents=True, exist_ok=True)
            output_path = Path(output_dir) / f"{Path(image_path).stem}_analysis.jpg"
            cv2.imwrite(str(output_path), result['visualization'])
            print(f"可视化结果已保存: {output_path}")

    else:
        print(f"分析失败: {result.get('error', '未知错误')}")


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 3:
        test_integrated_analyzer(sys.argv[1], sys.argv[2], sys.argv[3])
    else:
        print("用法: python integrated_analyzer.py <image_path> <yolo_model> <sam_model>")
