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) -> 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' / 'best.pt').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):
#     """
#     使用指定的模型对图像源进行预测。
#     """
#     if not Path(source).exists():
#         logging.error(f"错误：预测源路径不存在: {source}")
#         return

#     try:
#         logging.info(f"正在加载模型: {model_path}")
#         model = YOLO(model_path)
#         logging.info("模型加载成功。")

#         logging.info(f"正在对源进行预测: {source}")
#         results = model.predict(
#             source=source,
#             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,          # <<< ADDED: 保存与原图一样大的分割图
#         )
        
#         final_save_dir = Path(project_name) / "prediction"
#         logging.info(f"预测完成。结果保存在: {final_save_dir}")

#     except Exception as e:
#         logging.error(f"预测过程中发生错误: {e}")

def predict(model_path: str, source: str, project_name: str):
    """
    使用指定模型进行预测，并将所有类别的掩码合并到一个单通道灰度图中。
    在该图中，每个像素的值对应其类别ID (0=背景, 1=类别1, 2=类别2, ...)。
    """
    # --- 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}")
    except Exception as e:
        logging.error(f"加载或解析 dataset.yaml 失败: {e}")
        return

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

    try:
        # --- 2. 加载模型并执行预测 ---
        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
        )
        
        # --- 3. 创建新的输出目录并处理结果 ---
        mask_save_dir = Path(project_name) / "predicted_raw_masks"
        mask_save_dir.mkdir(parents=True, exist_ok=True)
        
        logging.info(f"预测完成。正在处理结果并保存组合掩码图于: {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} 中未检测到任何物体，将创建一张全黑的掩码图。")
                # 如果没有检测到物体，直接保存空白掩码图并继续处理下一张图片
                mask_filename = mask_save_dir / f"{original_image_path.stem}.png"
                cv2.imwrite(str(mask_filename), combined_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

            # YOLO结果通常按置信度从高到低排序。
            # 我们反向迭代（从低置信度到高置信度），以确保在掩码重叠区域，
            # 置信度更高的物体类别能够覆盖置信度较低的。
            for i in reversed(range(len(masks_data))):
                instance_mask = masks_data[i].cpu().numpy().astype(bool)
                class_id = class_ids[i]
                
                # 将掩码区域的像素值设置为其对应的类别ID。
                # 我们假设类别ID 0 是背景，所以即使模型预测了ID为0的物体，它也会被视为背景。
                if class_id != 0:
                    combined_mask[instance_mask] = class_id

            # --- 保存最终的组合掩码图 ---
            mask_filename = mask_save_dir / f"{original_image_path.stem}.png"
            cv2.imwrite(str(mask_filename), combined_mask)

        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="图片或图片文件夹的路径。"
    )
    args = parser.parse_args()
    
    # 2. 根据 --model 参数查找对应的训练历史
    available_runs = find_trained_models(config.PREDICT_BEST_MODEL_DIR, args.model)
    
    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' / 'best.pt'
        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
            )