import logging, sys, argparse, yaml
from pathlib import Path
from ultralytics import YOLO
import yolo_config as config
config.show_config_summary()  # 显示配置信息
from typing import List
import cv2
import numpy as np
# +----------------------------------------------------------------------------+
# ../BestMode_Predict_Results_DataSet_Public/
# └── YOLOv8n-seg_2025-09-20_10-00-00/  # The folder for the trained model you selected
#     ├── weights/
#     │   └── best.pt
#     │
#     ├── prediction/                    # Visualized results (image with colored mask overlay)
#     │   ├── image1.jpg
#     │   └── ...
#     │
#     └── prediction_masks_combined/     # <-- Your new unique folder for combined masks
#         ├── image1.png             # Grayscale mask, original size, pixel values are class IDs
#         ├── image2.png             # Grayscale mask
#         └── ...
# +----------------------------------------------------------------------------+

# --- 日志设置 ---
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[logging.StreamHandler()]
)

# --- 使用带 model_key 过滤的版本 ---
def find_trained_models(outputs_dir: Path, model_key: str, pt_name: str) -> List[str]:
    """
    扫描输出目录，查找特定基础模型的所有有效的、已完成的训练项目。
    """
    trained_models = []
    if not outputs_dir.is_dir():
        logging.warning(f"输出目录不存在: {outputs_dir}")
        return []
    
    for project_folder in outputs_dir.iterdir():
        if project_folder.is_dir() and project_folder.name.startswith(model_key + '_'):
            if (project_folder / 'weights' / pt_name).exists():
                trained_models.append(project_folder.name)
    
    trained_models.sort()
    return trained_models

# 老版predict函数，只有默认输出
# def predict_old(model_path: str, source: str, project_name: str):
#     """
#     使用指定模型进行预测。
#     1. 保存原始灰度掩码 (predicted_raw_masks)，像素值为类别ID。
#     2. 保存彩色可视化掩码 (predicted_color_masks)，根据 color_map.yaml 上色。
#     """
    
#     # --- 1. 加载类别名称 (用于日志记录和确定类别数量) ---
#     try:
#         with open(config.DATASET_YAML_PATH, 'r', encoding='utf-8') as f:
#             class_names = yaml.safe_load(f)['names']
#         logging.info(f"成功从 {config.DATASET_YAML_PATH} 加载类别名称: {class_names}")
#         num_classes = len(class_names)
#         logging.info(f"检测到 {num_classes} 个类别 (ID 0 到 {num_classes-1})。")
#     except Exception as e:
#         logging.error(f"加载或解析 dataset.yaml 失败: {e}")
#         return

#     # --- [新增] 2. 加载颜色查找表 (LUT) ---
#     color_lut = np.zeros((num_classes, 3), dtype=np.uint8) # 默认全黑
#     color_map_path = 'color_map.yaml'
    
#     try:
#         with open(color_map_path, 'r', encoding='utf-8') as f:
#             color_data = yaml.safe_load(f)['colors']
#         logging.info(f"正在从 {color_map_path} 加载颜色...")
        
#         for class_id, bgr_value in color_data.items():
#             if 0 <= class_id < num_classes:
#                 color_lut[class_id] = bgr_value
#             else:
#                 logging.warning(f"color_map.yaml 中的类别ID {class_id} 超出范围 (0-{num_classes-1})，已忽略。")
#         logging.info("颜色查找表 (LUT) 创建成功。")
        
#     except FileNotFoundError:
#         logging.error(f"错误：未找到 {color_map_path}。")
#         logging.warning("将使用随机颜色作为备用方案（背景除外）。")
#         # 创建随机颜色作为备用
#         for i in range(1, num_classes): # 保持 0 (背景) 为黑色 [0,0,0]
#             color_lut[i] = np.random.randint(0, 255, 3, dtype=np.uint8)
#     except Exception as e:
#         logging.error(f"加载或解析 {color_map_path} 失败: {e}", exc_info=True)
#         return

#     if not Path(source).exists():
#         logging.error(f"错误：预测源路径不存在: {source}")
#         return

#     try:
#         # --- 3. 加载模型并执行预测 ---
#         logging.info(f"正在加载模型: {model_path}")
#         model = YOLO(model_path)
#         logging.info("模型加载成功。")

#         logging.info(f"正在对源进行预测: {source}")
#         results = model.predict(
#             source=source,
#             stream=True, 
#             save=config.SAVE_PREDICTIONS,
#             show_labels=config.SHOW_LABELS,
#             show_conf=config.SHOW_CONF,
#             project=project_name,
#             name="prediction",
#             exist_ok=True,
#             retina_masks=True,
#             conf=0.1
#         )
        
#         # --- 4. 创建新的输出目录并处理结果 ---
#         # [修改] 定义两个保存目录
#         raw_mask_save_dir = Path(project_name) / "predicted_raw_masks"
#         color_mask_save_dir = Path(project_name) / "predicted_color_masks" # [新增] 彩色图目录
        
#         raw_mask_save_dir.mkdir(parents=True, exist_ok=True)
#         color_mask_save_dir.mkdir(parents=True, exist_ok=True) # [新增]
        
#         # [修改] 更新日志信息
#         logging.info(f"预测完成。正在处理结果...")
#         logging.info(f"  -> 原始灰度掩码将保存于: {raw_mask_save_dir}")
#         logging.info(f"  -> 彩色可视化掩码将保存于: {color_mask_save_dir}")

#         for result in results:
#             original_image_path = Path(result.path)
#             logging.info(f"正在处理图片: {original_image_path.name}")

#             h, w = result.orig_shape
#             # 创建一个空白画布，用于存储所有类别的组合掩码。0 代表背景。
#             combined_mask = np.zeros((h, w), dtype=np.uint8)

#             if result.masks is None:
#                 logging.warning(f"  -> 在图片 {original_image_path.name} 中未检测到任何物体，将创建一张全黑的掩码图。")
                
#                 # --- 保存原始灰度掩码 (全黑) ---
#                 raw_mask_filename = raw_mask_save_dir / f"{original_image_path.stem}.png"
#                 cv2.imwrite(str(raw_mask_filename), combined_mask)
                
#                 # --- [新增] 保存彩色掩码 (全黑) ---
#                 # 即使是全黑的，也应用LUT，结果还是全黑
#                 colorized_mask = color_lut[combined_mask] # (H, W, 3)
#                 color_mask_filename = color_mask_save_dir / f"{original_image_path.stem}.png"
#                 cv2.imwrite(str(color_mask_filename), colorized_mask)

#                 continue # 继续处理下一张图片

#             # --- 核心：创建组合掩码 ---
#             masks_data = result.masks.data
#             class_ids = result.boxes.cls.int().cpu().numpy()

#             if len(masks_data) != len(class_ids):
#                  logging.error(f"  -> 掩码和类别ID数量不匹配，跳过图片 {original_image_path.name}")
#                  continue

#             # (核心逻辑保持不变)
#             # 反向迭代，确保高置信度覆盖低置信度
#             for i in reversed(range(len(masks_data))):
#                 instance_mask = masks_data[i].cpu().numpy().astype(bool)
#                 class_id = class_ids[i]
                
#                 if class_id != 0: # 假设 0 是背景
#                     combined_mask[instance_mask] = class_id

#             # --- 保存 1: 最终的组合掩码图 (灰度) ---
#             raw_mask_filename = raw_mask_save_dir / f"{original_image_path.stem}.png"
#             cv2.imwrite(str(raw_mask_filename), combined_mask)

#             # --- [新增] 保存 2: 最终的彩色掩码图 (RGB) ---
#             # 这是最高效的上色方法：
#             # color_lut (14, 3)
#             # combined_mask (H, W)，值 0-13
#             # NumPy 会自动将 (H, W) 中的每个值作为索引去 color_lut 中取 (B,G,R) 元组
#             colorized_mask = color_lut[combined_mask] # 结果维度 (H, W, 3)
            
#             color_mask_filename = color_mask_save_dir / f"{original_image_path.stem}.png"
#             cv2.imwrite(str(color_mask_filename), colorized_mask)

#         logging.info(f"所有掩码图已成功保存。")

#     except Exception as e:
#         logging.error(f"预测或手动保存过程中发生错误: {e}", exc_info=True)

def predict(model_path: str, source: str, project_name: str, pt_name: str, conf_threshold: float):
    """
    使用指定模型进行预测。
    1. 保存原始灰度掩码 (predicted_raw_masks)，像素值为类别ID。
    2. 保存彩色可视化掩码 (predicted_color_masks)。
    3. 保存三合一对比图 (predicted_comparison) [Original | Ground Truth | Prediction]。
    """
    # 
    pt_name_raw = pt_name.replace('.pt','') # pt_name去掉.pt后缀，方便文件命名使用
    # [新增] 标志，用于决定是否将(类别0..N-1) 偏移到 (1..N)，为 "真背景" 腾出 0
    perform_id_shift = False
    
    # --- 1. 加载类别名称和颜色 ---
    try:
        with open(config.DATASET_YAML_PATH, 'r', encoding='utf-8') as f:
            yaml_data = yaml.safe_load(f)
            class_names = yaml_data.get('names')
            color_data_from_dataset = yaml_data.get('colors')
        
        if not class_names:
            logging.error(f"在 {config.DATASET_YAML_PATH} 中未找到 'names' 键。")
            return
            
        logging.info(f"成功从 {config.DATASET_YAML_PATH} 加载类别名称: {class_names}")
        num_classes = len(class_names)
        logging.info(f"检测到 {num_classes} 个类别 (ID 0 到 {num_classes-1})。")

    except Exception as e:
        logging.error(f"加载或解析 dataset.yaml 失败: {e}")
        return

    # --- 2. 创建颜色查找表 (LUT) ---
    # [调整] 先创建一个临时的 (N) 维 LUT
    temp_lut = np.zeros((num_classes, 3), dtype=np.uint8) 

    if color_data_from_dataset:
        logging.info(f"正在从 {config.DATASET_YAML_PATH} 加载颜色...")
        try:
            for class_id, rgb_value in color_data_from_dataset.items():
                if 0 <= class_id < num_classes:
                    temp_lut[class_id] = rgb_value[::-1] # BGR
            logging.info("临时颜色查找表 (LUT) 已成功从 dataset.yaml 创建。")
        except Exception as e:
            logging.error(f"解析 'colors' 键失败: {e}。将尝试备用方案。")
            color_data_from_dataset = None # 强制使用备用方案
            
    if not color_data_from_dataset:
        # 备用方案 1: 尝试 color_map.yaml
        color_map_path = 'color_map.yaml'
        logging.warning(f"未从 dataset.yaml 加载颜色，正在尝试备用方案: {color_map_path}")
        try:
            with open(color_map_path, 'r', encoding='utf-8') as f:
                color_data = yaml.safe_load(f)['colors']
            for class_id, bgr_value in color_data.items():
                if 0 <= class_id < num_classes:
                    color_lut[class_id] = bgr_value
            logging.info(f"颜色查找表 (LUT) 已成功从 {color_map_path} 创建。")
        except Exception:
            # 备用方案 2: 随机颜色
            logging.warning(f"{color_map_path} 未找到或解析失败。将使用随机颜色。")
            for i in range(1, num_classes): # 保持 0 (背景) 为黑色 [0,0,0]
                color_lut[i] = np.random.randint(0, 255, 3, dtype=np.uint8)

    # --- [新增] 核心条件判断 ---
    # 检查 类别0 (索引0) 的颜色是否为 BGR(0,0,0)
    if np.any(temp_lut[0] != [0, 0, 0]):
        # 颜色不是黑色，执行 "ID+1" 偏移
        logging.info("检测到 类别0 的颜色不是黑色。将执行 ID+1 偏移，使 0 成为专用背景。")
        perform_id_shift = True
        # [调整] 创建一个 (N+1) 维的最终 LUT
        color_lut = np.zeros((num_classes + 1, 3), dtype=np.uint8)
        # 将原 0..N-1 的颜色 复制到 1..N
        color_lut[1:] = temp_lut
        # 索引 0 保持为 [0, 0, 0] (真背景)
    else:
        # 颜色是黑色，不执行偏移
        logging.info("检测到 类别0 的颜色是黑色。将 类别0 视为背景。")
        perform_id_shift = False
        color_lut = temp_lut # [调整] 使用 (N) 维的原始 LUT

    if not Path(source).exists():
        logging.error(f"错误：预测源路径不存在: {source}")
        return

    try:
        # --- 3. 加载模型并执行预测 ---
        logging.info(f"正在加载模型: {model_path}")
        model = YOLO(model_path)
        logging.info("模型加载成功。")

        logging.info(f"正在对源进行预测: {source}")
        results = model.predict(
            source=source,
            stream=True, 
            save=False,             # <-- 关闭保存，之后手动处理
            show_labels=config.SHOW_LABELS,
            show_conf=config.SHOW_CONF,
            # project=project_name, # <-- 不再需要，我们手动处理
            # name="prediction",    # <-- 不再需要，我们手动处理
            exist_ok=True,
            retina_masks=True,
            conf=conf_threshold # TODO
        )

        # --- 4. [扩展] 创建所有输出目录 ---
        yolo_prediction_dir = Path(project_name) / "prediction"
        raw_mask_save_dir = Path(project_name) / "predicted_raw_masks"
        color_mask_save_dir = Path(project_name) / "predicted_color_masks"
        comparison_save_dir = Path(project_name) / "predicted_comparison" # [新增]

        logging.info(f"预测完成。正在处理结果...")
        logging.info(f"  -> Yolo预测结果将保存于: {yolo_prediction_dir}")
        logging.info(f"  -> 原始灰度掩码将保存于: {raw_mask_save_dir}")
        logging.info(f"  -> 彩色预测掩码将保存于: {color_mask_save_dir}")
        logging.info(f"  -> 三合一对比图将保存于: {comparison_save_dir}") # [新增]

        yolo_prediction_dir.mkdir(parents=True, exist_ok=True)
        raw_mask_save_dir.mkdir(parents=True, exist_ok=True)
        color_mask_save_dir.mkdir(parents=True, exist_ok=True)
        comparison_save_dir.mkdir(parents=True, exist_ok=True) # [新增]
        
        # --- 辅助函数：在图像上添加标签 ---
        def add_label_to_image(img, label):
            """在图像左上角添加一个带背景的标签"""
            img_with_label = img.copy()
            h, w = img_with_label.shape[:2]
            # 动态调整字体大小和粗细
            font_scale = max(0.8, w // 1000)
            thickness = max(1, w // 500)
            
            (text_w, text_h), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness)
            
            # 绘制黑色背景矩形
            cv2.rectangle(img_with_label, (0, 0), (text_w + 20, text_h + 30), (0, 0, 0), -1)
            # 绘制白色文字
            cv2.putText(img_with_label, label, (10, text_h + 15), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), thickness, cv2.LINE_AA)
            return img_with_label
        
        # --- 5. [扩展] 处理结果循环 ---
        for result in results:
            original_image_path = Path(result.path)
            logging.info(f"正在处理图片: {original_image_path.name}")

            h, w = result.orig_shape

            # --- 5a. 保存 Yolo-predict 本身结果 ---
            if config.SAVE_PREDICTIONS == True:
                # a. 获取原始图片的文件名和后缀
                # result.path 是原始图片的完整路径
                original_path = Path(result.path)
                image_stem = original_path.stem    # 这是 "图片名"
                image_suffix = original_path.suffix  # 这是 ".后缀" (例如 .jpg)
                # b. 构建你的新文件名
                # 格式: 图片名_{pt_name_raw}.后缀
                new_filename = f"{image_stem}_{pt_name_raw}{image_suffix}"
                # c. 构建完整的保存路径
                save_path = yolo_prediction_dir / new_filename
                # d. 获取绘制了检测框的图像 (NumPy 数组)
                # result.plot() 会自动使用你在 predict 中设置的 show_labels, show_conf 等参数
                annotated_image = result.plot()
                # e. 使用 OpenCV (cv2) 保存图像
                cv2.imwrite(str(save_path), annotated_image)
                print(f"已保存: {save_path}")

            # --- 5b. [新增] 查找并加载 Ground Truth 掩码 ---
            gt_mask = None
            try:
                # 假设的路径结构: .../DataSet_Public/5_.../images/val/image.png
                # 我们要找: .../DataSet_Public/5_.../labels_GT/val/image.png
                image_filename = original_image_path.name
                image_parent_dir_name = original_image_path.parent.name # 'val' or 'train'
                images_dir = original_image_path.parent.parent # '.../images'
                dataset_root_dir = images_dir.parent # '.../5_My_Gastric_2025_10_29'
                
                gt_path = dataset_root_dir / "labels_GT" / image_parent_dir_name / image_filename
                
                if not gt_path.exists():
                    logging.warning(f"  -> 未找到 Ground Truth 文件: {gt_path}")
                else:
                    gt_mask = cv2.imread(str(gt_path), cv2.IMREAD_GRAYSCALE)
                    if gt_mask is None:
                         logging.warning(f"  -> 读取 Ground Truth 文件失败: {gt_path}")
                    elif gt_mask.shape != (h, w):
                         logging.warning(f"  -> GT 掩码形状 {gt_mask.shape} 与图像形状 {(h, w)} 不匹配。正在调整GT大小...")
                         gt_mask = cv2.resize(gt_mask, (w, h), interpolation=cv2.INTER_NEAREST)
            except Exception as e:
                logging.error(f"  -> 查找或加载GT文件时出错: {e}", exc_info=True)
                gt_mask = None

            # --- 5c. 创建预测掩码 (与之前相同) ---
            combined_mask = np.zeros((h, w), dtype=np.uint8)

            if result.masks is None:
                logging.warning(f"  -> 在图片 {original_image_path.name} 中未检测到任何物体。")
            else:
                masks_data = result.masks.data
                class_ids = result.boxes.cls.int().cpu().numpy()

                if len(masks_data) != len(class_ids):
                     logging.error(f"  -> 掩码和类别ID数量不匹配，跳过图片 {original_image_path.name}")
                     continue

                for i in reversed(range(len(masks_data))):
                    instance_mask = masks_data[i].cpu().numpy().astype(bool)
                    class_id = class_ids[i]
                    # [调整] 根据标志执行不同逻辑
                    if perform_id_shift:
                        # 偏移模式：将所有 ID (包括0) 都 +1
                        combined_mask[instance_mask] = class_id + 1
                    else:
                        # 原始模式：忽略 ID 0
                        if class_id != 0:
                            combined_mask[instance_mask] = class_id

            # --- 5d. [扩展] 保存所有三种输出 ---

            # --- 保存 1: 原始灰度掩码 (预测) ---
            raw_mask_filename = raw_mask_save_dir / f"{original_image_path.stem}_{pt_name_raw}.png"
            cv2.imwrite(str(raw_mask_filename), combined_mask)

            # --- 保存 2: 彩色预测掩码 ---
            colorized_pred_mask = color_lut[combined_mask] # (H, W, 3)
            color_mask_filename = color_mask_save_dir / f"{original_image_path.stem}_{pt_name_raw}.png"
            cv2.imwrite(str(color_mask_filename), colorized_pred_mask)

            # --- [新增] 保存 3: 三合一对比图 ---
            try:
                # 1. 获取原始图
                original_image_bgr = result.orig_img # YOLO 结果中自带 BGR 格式原图
                
                # 2. 获取彩色 GT 图
                if gt_mask is not None:
                    colorized_gt_mask = color_lut[gt_mask]
                else:
                    # 如果GT不存在，创建一个黑色占位图
                    logging.warning(f"  -> 在对比图中将使用黑色图像作为 Ground Truth 占位符。")
                    colorized_gt_mask = np.zeros_like(original_image_bgr) # (H, W, 3)

                # 3. 获取彩色预测图 (已在上面生成: colorized_pred_mask)

                # 确保所有图像形状一致 (理论上应该一致)
                if not (original_image_bgr.shape == colorized_gt_mask.shape == colorized_pred_mask.shape):
                     logging.error(f"  -> 形状不匹配! 原图: {original_image_bgr.shape}, GT: {colorized_gt_mask.shape}, 预测: {colorized_pred_mask.shape}。跳过对比图。")
                     continue

                # 为每张图添加标签
                img_labeled = add_label_to_image(original_image_bgr, "Original")
                gt_labeled = add_label_to_image(colorized_gt_mask, "Ground Truth")
                pred_labeled = add_label_to_image(colorized_pred_mask, "Prediction")

                # 水平拼接
                comparison_image = cv2.hconcat([img_labeled, gt_labeled, pred_labeled])
                
                # 保存 (使用 .jpg 格式以节省空间)
                comparison_filename = comparison_save_dir / f"{original_image_path.stem}_{pt_name_raw}.jpg"
                cv2.imwrite(str(comparison_filename), comparison_image)

            except Exception as e:
                logging.error(f"  -> 创建或保存对比图失败: {e}", exc_info=True)


        logging.info(f"所有掩码图和对比图已成功保存。")

    except Exception as e:
        logging.error(f"预测或手动保存过程中发生错误: {e}", exc_info=True)

if __name__ == "__main__":
    # 1. 创建解析器，现在只需要 --model 和 --source
    parser = argparse.ArgumentParser(description="使用已训练的YOLO模型进行预测。")
    
    parser.add_argument(
        "--model",
        type=str,
        required=True,
        choices=list(config.MODEL_CONFIGS.keys()),
        help="选择一个基础模型类型，以筛选其训练历史。"
    )
    parser.add_argument(
        "--source",
        type=str,
        default=str(config.TEST_IMAGE_DIR),
        help="图片或图片文件夹的路径。"
    )
    parser.add_argument(
        "--pt_name",
        type=str,
        default=str("best.pt"),
        help="图片或图片文件夹的路径。"
    )
    parser.add_argument(
        "--conf",
        type=float,  # 1. 类型应为浮点数
        default=0.2,
        help="设置预测的置信度阈值 (例如: 0.25)"  # 2. 补充 help 信息
    )
    args = parser.parse_args()
    
    # 2. 根据 --model 参数查找对应的训练历史
    available_runs = find_trained_models(config.PREDICT_BEST_MODEL_DIR, args.model, args.pt_name)
    
    run_to_use = None

    # 3. 根据找到的结果数量，决定下一步操作
    if not available_runs:
        logging.error(f"错误：在 {config.PREDICT_BEST_MODEL_DIR} 中没有找到任何关于模型 '{args.model}' 的有效训练记录。")
        sys.exit(1)
    
    elif len(available_runs) == 1:
        # 如果只有一个结果，自动选择
        run_to_use = available_runs[0]
        logging.info(f"只找到一个训练版本，已自动选择: {run_to_use}")
    
    else:
        # 如果有多个结果，进入交互式选择模式
        print(f"\n为模型 '{args.model}' 找到多个训练版本:")
        for i, run_name in enumerate(available_runs, 1):
            print(f"  [{i}] {run_name}")
        
        while True:
            try:
                choice = input(f"请输入您想使用的版本序号 (1-{len(available_runs)}): ")
                choice_index = int(choice)
                if 1 <= choice_index <= len(available_runs):
                    run_to_use = available_runs[choice_index - 1]
                    break
                else:
                    print("错误：输入无效，请输入列表中的序号。")
            except ValueError:
                print("错误：请输入一个数字。")
            except (KeyboardInterrupt, EOFError):
                print("\n操作已取消。")
                sys.exit(0)

    # 4. 使用最终确定的 run_to_use 来执行预测
    if run_to_use:
        model_to_use = config.PREDICT_BEST_MODEL_DIR / run_to_use / 'weights' / args.pt_name
        project_to_save_in = str(config.PREDICT_BEST_MODEL_DIR / run_to_use)

        if not model_to_use.exists():
            logging.error(f"严重错误：找不到权重文件 {model_to_use}。")
        else:
            predict(
                model_path=str(model_to_use), 
                source=args.source, 
                project_name=project_to_save_in,
                pt_name = args.pt_name,
                conf_threshold=args.conf
            )