"""
香菇检测与分割模块

集成YOLO检测和SAM分割功能
"""

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

try:
    from ultralytics import YOLO
    YOLO_AVAILABLE = True
except ImportError:
    YOLO_AVAILABLE = False
    logging.warning("ultralytics未安装，请运行：pip install ultralytics")

try:
    from segment_anything import sam_model_registry, SamPredictor
    SAM_AVAILABLE = True
except ImportError:
    SAM_AVAILABLE = False
    logging.warning("segment_anything未安装，请运行：pip install git+https://github.com/facebookresearch/segment-anything.git")


class MushroomDetector:
    """香菇检测器（YOLO）"""

    def __init__(self, model_path: str, confidence_threshold: float = 0.5):
        """
        初始化检测器

        参数:
            model_path: YOLO模型路径
            confidence_threshold: 置信度阈值
        """
        if not YOLO_AVAILABLE:
            raise ImportError("YOLO未安装")

        self.model = YOLO(model_path)
        self.confidence_threshold = confidence_threshold
        logging.info(f"YOLO检测器加载完成：{model_path}")

    def detect(self, image: np.ndarray) -> List[Dict]:
        """
        检测图像中的香菇

        参数:
            image: BGR格式图像

        返回:
            检测结果列表，每个元素包含：
            - box: [x1, y1, x2, y2] 边界框坐标
            - confidence: 置信度
            - class_id: 类别ID
        """
        results = self.model(image, verbose=False)

        detections = []
        for result in results:
            boxes = result.boxes
            for box in boxes:
                conf = float(box.conf[0])
                if conf >= self.confidence_threshold:
                    xyxy = box.xyxy[0].cpu().numpy()
                    cls = int(box.cls[0])

                    detections.append({
                        'box': xyxy.tolist(),
                        'confidence': conf,
                        'class_id': cls
                    })

        logging.info(f"检测到 {len(detections)} 个香菇")
        return detections

    def crop_roi(self, image: np.ndarray, box: List[float],
                 padding_ratio: float = 0.1) -> Tuple[np.ndarray, Tuple[int, int]]:
        """
        根据检测框裁剪ROI区域（带padding）

        参数:
            image: 原始图像
            box: [x1, y1, x2, y2] 边界框
            padding_ratio: 扩展比例

        返回:
            (裁剪后的图像, (offset_x, offset_y))
        """
        h, w = image.shape[:2]
        x1, y1, x2, y2 = map(int, box)

        # 计算padding
        box_w = x2 - x1
        box_h = y2 - y1
        pad_w = int(box_w * padding_ratio)
        pad_h = int(box_h * padding_ratio)

        # 扩展边界框
        x1_pad = max(0, x1 - pad_w)
        y1_pad = max(0, y1 - pad_h)
        x2_pad = min(w, x2 + pad_w)
        y2_pad = min(h, y2 + pad_h)

        # 裁剪
        roi = image[y1_pad:y2_pad, x1_pad:x2_pad]

        return roi, (x1_pad, y1_pad)


class MushroomSegmenter:
    """菌盖分割器（SAM）"""

    def __init__(self, model_path: str, model_type: str = 'vit_b', device: str = 'cuda'):
        """
        初始化SAM分割器

        参数:
            model_path: SAM模型权重路径
            model_type: 模型类型 (vit_b, vit_l, vit_h)
            device: 运行设备
        """
        if not SAM_AVAILABLE:
            raise ImportError("segment_anything未安装")

        # 检查设备
        if device == 'cuda' and not torch.cuda.is_available():
            device = 'cpu'
            logging.warning("CUDA不可用，使用CPU")

        self.device = device

        # 加载SAM模型
        sam = sam_model_registry[model_type](checkpoint=model_path)
        sam.to(device=device)

        self.predictor = SamPredictor(sam)
        logging.info(f"SAM分割器加载完成：{model_path} ({model_type}, {device})")

    def segment(self, image: np.ndarray, box: List[float]) -> Optional[np.ndarray]:
        """
        使用SAM分割菌盖

        参数:
            image: BGR格式图像
            box: [x1, y1, x2, y2] 边界框（作为Box Prompt）

        返回:
            二值掩码 (H, W)，菌盖区域为255，背景为0
            如果分割失败返回None
        """
        # 转换为RGB
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # 设置图像
        self.predictor.set_image(image_rgb)

        # 准备Box Prompt
        input_box = np.array(box)

        # 在检测框中心添加前景点提示，提高分割精度
        cx = (box[0] + box[2]) / 2
        cy = (box[1] + box[3]) / 2
        input_point = np.array([[cx, cy]])
        input_label = np.array([1])  # 1=前景

        # 预测（同时使用box和point prompt）
        masks, scores, _ = self.predictor.predict(
            point_coords=input_point,
            point_labels=input_label,
            box=input_box,
            multimask_output=True
        )

        # 选择得分最高的掩码
        best_mask_idx = np.argmax(scores)
        mask = masks[best_mask_idx]

        # 转换为uint8
        mask_uint8 = (mask * 255).astype(np.uint8)

        # 形态学后处理：平滑边缘
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
        mask_uint8 = cv2.morphologyEx(mask_uint8, cv2.MORPH_CLOSE, kernel)
        mask_uint8 = cv2.morphologyEx(mask_uint8, cv2.MORPH_OPEN, kernel)

        logging.info(f"分割完成，掩码得分: {scores[best_mask_idx]:.3f}")

        return mask_uint8

    def segment_roi(self, roi_image: np.ndarray,
                    original_box: List[float],
                    offset: Tuple[int, int]) -> Optional[np.ndarray]:
        """
        在ROI区域内分割菌盖（推荐使用，高效）

        这是ROI处理哲学的核心方法：
        1. 只处理裁剪后的小图（例如500×500而非4000×3000）
        2. 显著减少SAM的计算量（快10-20倍）
        3. 内存占用更小

        参数:
            roi_image: 裁剪后的ROI图像
            original_box: 原图中的检测框
            offset: ROI相对于原图的偏移量 (offset_x, offset_y)

        返回:
            ROI尺寸的二值掩码
        """
        # 将原图坐标转换为ROI坐标
        x1, y1, x2, y2 = original_box
        offset_x, offset_y = offset

        roi_box = [
            x1 - offset_x,
            y1 - offset_y,
            x2 - offset_x,
            y2 - offset_y
        ]

        # 只在ROI上运行SAM（关键优化点）
        return self.segment(roi_image, roi_box)


def test_detection_and_segmentation(
    image_path: str,
    yolo_model_path: str,
    sam_model_path: str
):
    """
    测试检测和分割功能

    参数:
        image_path: 测试图像路径
        yolo_model_path: YOLO模型路径
        sam_model_path: SAM模型路径
    """
    logging.basicConfig(level=logging.INFO)

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

    # 创建检测器
    detector = MushroomDetector(yolo_model_path)

    # 检测香菇
    detections = detector.detect(image)

    if len(detections) == 0:
        print("未检测到香菇")
        return

    print(f"\n检测到 {len(detections)} 个香菇")

    # 创建分割器
    segmenter = MushroomSegmenter(sam_model_path)

    # 对第一个检测结果进行分割
    detection = detections[0]
    box = detection['box']

    print(f"处理第一个香菇，置信度: {detection['confidence']:.3f}")

    # 分割
    mask = segmenter.segment(image, box)

    if mask is not None:
        # 可视化
        vis_image = image.copy()

        # 绘制检测框
        x1, y1, x2, y2 = map(int, box)
        cv2.rectangle(vis_image, (x1, y1), (x2, y2), (0, 255, 0), 2)

        # 叠加掩码
        mask_color = np.zeros_like(image)
        mask_color[mask > 0] = [0, 0, 255]
        vis_image = cv2.addWeighted(vis_image, 0.7, mask_color, 0.3, 0)

        # 保存结果
        output_path = image_path.replace('.', '_detected_segmented.')
        cv2.imwrite(output_path, vis_image)
        print(f"结果已保存到：{output_path}")
    else:
        print("分割失败")


if __name__ == "__main__":
    import sys

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