import cv2
import numpy as np
import yaml
import logging
import argparse
import sys
from pathlib import Path
from tqdm import tqdm

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

try:
    import yolo_config as config
    logging.info(f"成功加载配置 yolo_config.py")
    logging.info(f"  -> 最佳模型目录: {config.PREDICT_BEST_MODEL_DIR}")
    logging.info(f"  -> 数据集 YAML: {config.DATASET_YAML_PATH}")
    logging.info(f"  -> 测试图片目录: {config.TEST_IMAGE_DIR}")
except ImportError:
    logging.error("错误: yolo_config.py 未找到。请确保它在同一目录下。")
    sys.exit(1)
except AttributeError as e:
    logging.error(f"错误: yolo_config.py 中缺少必要的配置。 {e}")
    sys.exit(1)

# --- 辅助函数: 1. 添加标签 ---
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

# --- 辅助函数: 2. 加载颜色LUT ---
def load_color_lut(dataset_yaml_path):
    """
    加载颜色查找表 (LUT)，精确复制 yolo_predict_V2.py 的逻辑
    """
    try:
        with open(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"在 {dataset_yaml_path} 中未找到 'names' 键。")
            return None
            
        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 None

    # 创建 (N) 维临时 LUT
    temp_lut = np.zeros((num_classes, 3), dtype=np.uint8) 

    if color_data_from_dataset:
        logging.info(f"正在从 {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:
        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:
                    temp_lut[class_id] = bgr_value # 假设 color_map.yaml 存的是 BGR
            logging.info(f"颜色查找表 (LUT) 已成功从 {color_map_path} 创建。")
        except Exception:
            logging.warning(f"{color_map_path} 未找到或解析失败。将使用随机颜色。")
            for i in range(1, num_classes):
                temp_lut[i] = np.random.randint(0, 255, 3, dtype=np.uint8)

    # --- 核心条件判断 (与 yolo_predict_V2.py 完全一致) ---
    if np.any(temp_lut[0] != [0, 0, 0]):
        logging.info("检测到 类别0 的颜色不是黑色。将创建 (N+1) LUT，使 0 成为专用背景。")
        color_lut = np.zeros((num_classes + 1, 3), dtype=np.uint8)
        color_lut[1:] = temp_lut
    else:
        logging.info("检测到 类别0 的颜色是黑色。将使用 (N) LUT。")
        color_lut = temp_lut

    logging.info(f"最终颜色LUT创建成功，维度: {color_lut.shape}")
    return color_lut

# --- 辅助函数: 3. 查找模型目录 ---
def find_model_dirs(base_dir, required_subdir):
    """查找所有包含 'required_subdir' 子目录的有效模型文件夹"""
    model_dirs = []
    if not base_dir.is_dir():
        logging.error(f"基础目录不存在: {base_dir}")
        return []
        
    for d in base_dir.iterdir():
        # 必须是一个目录，并且包含所需的预测结果子目录
        if d.is_dir() and (d / required_subdir).exists():
            model_dirs.append(d)
            
    model_dirs.sort()
    return model_dirs

# --- 辅助函数: 4. 查找图像名称 ---
def find_image_names(model_dirs, sub_dir_name, file_glob_patterns):
    """
    从第一个模型目录的指定子目录(sub_dir_name)中获取所有待处理的图像文件名
    使用 file_glob_patterns 列表 (例如 ['*.jpg', '*.png'])
    """
    image_names = set() # 使用集合避免重复
    if not model_dirs:
        return []
    
    first_model_target_dir = model_dirs[0] / sub_dir_name
    
    for pattern in file_glob_patterns:
        for f in first_model_target_dir.glob(pattern):
            if f.is_file():
                image_names.add(f.name)
            
    sorted_names = sorted(list(image_names))
    return sorted_names

# --- 辅助函数: 5. 获取原图和GT路径 ---
def get_gt_and_original_paths(image_name):
    """
    根据输出的掩码文件名，反向推导原始图像和GT图像的路径。
    假设文件名格式为: {original_stem}_{pt_name_raw}.png
    """
    try:
        # 1. 从 config 中获取关键路径
        test_image_dir = config.TEST_IMAGE_DIR
        images_dir = test_image_dir.parent
        dataset_root_dir = images_dir.parent
        image_parent_dir_name = test_image_dir.name # e.g., 'val'

        # 2. 推导原始图像的 stem
        if '_' not in image_name:
             logging.warning(f"  -> 图像 {image_name} 似乎不含 '_{{pt_name}}' 后缀，将尝试使用完整文件名作为 stem。")
             original_stem = Path(image_name).stem
        else:
             # 从最后一个 '_' 拆分
             original_stem = image_name[:image_name.rfind('_')]


        # 3. 在 test_image_dir 中查找原始图像
        original_image_path = None
        # glob 查找, 匹配 .png, .jpg, .bmp 等
        for f in test_image_dir.glob(f"{original_stem}.*"): 
            if f.is_file() and f.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff']:
                original_image_path = f
                break # 找到第一个匹配项
        
        if not original_image_path:
            # 尝试在上一级目录查找 (兼容某些奇怪的 val/test 结构)
            for f in test_image_dir.parent.glob(f"{original_stem}.*"):
                 if f.is_file() and f.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff']:
                    original_image_path = f
                    break
        
        if not original_image_path:
            logging.warning(f"  -> 在 {test_image_dir} (及其父目录) 中未找到原始图像 (如 {original_stem}.*)")
            return None, None
            
        # 4. 推导GT路径 (逻辑复制自 yolo_predict_V2.py)
        image_filename = original_image_path.name
        
        gt_path = dataset_root_dir / "labels_GT" / image_parent_dir_name / image_filename
        gt_path_png = gt_path.with_suffix('.png') # 备用

        if gt_path.exists():
             return original_image_path, gt_path
        elif gt_path_png.exists():
             return original_image_path, gt_path_png
        else:
             # 返回预期的路径，主循环将处理 "文件未找到"
             return original_image_path, gt_path
             
    except Exception as e:
        logging.error(f"  -> 推导路径时出错: {e}", exc_info=True)
        return None, None


# --- [新增] 核心处理函数 ---
def generate_comparison_images(base_dir, output_dir_name, source_subdir_name, 
                                 gt_panel_label, use_gt_overlay, 
                                 file_glob_patterns, color_lut):
    """
    核心处理函数，用于生成一种类型的对比图。
    """
    logging.info(f"--- 🚀 开始生成 '{output_dir_name}' ---")
    
    output_dir = base_dir / output_dir_name
    output_dir.mkdir(exist_ok=True)
    
    # 3. 查找模型目录
    model_dirs = find_model_dirs(base_dir, source_subdir_name)
    if not model_dirs:
        logging.error(f"在 {base_dir} 中未找到任何包含 '{source_subdir_name}' 的有效模型预测目录。")
        logging.warning(f"--- ⚠️ 跳过 '{output_dir_name}' 的生成 ---")
        return
    logging.info(f"找到 {len(model_dirs)} 个模型进行对比:")
    for d in model_dirs:
        logging.info(f"  -> {d.name}")

    # 4. 查找图像列表
    image_names = find_image_names(model_dirs, source_subdir_name, file_glob_patterns)
    if not image_names:
        logging.error(f"未找到任何匹配 {file_glob_patterns} 的图像文件。")
        logging.error(f"请检查 {model_dirs[0] / source_subdir_name} 目录。")
        logging.warning(f"--- ⚠️ 跳过 '{output_dir_name}' 的生成 ---")
        return
    logging.info(f"找到 {len(image_names)} 张图像进行处理。")

    # 5. --- 遍历每张图像 ---
    for image_name in tqdm(image_names, desc=f"生成 {output_dir_name}"):
        
        comparison_panels = []
        
        # 5a. 获取路径
        orig_path, gt_path = get_gt_and_original_paths(image_name)
        
        if not orig_path:
            logging.warning(f"\n跳过: 未能找到 {image_name} 的原始图像。")
            continue
            
        # 5b. 加载原图
        orig_img = cv2.imread(str(orig_path))
        if orig_img is None:
            logging.warning(f"\n跳过: 无法读取原始图像 {orig_path}")
            continue
            
        h, w = orig_img.shape[:2]
        
        # 1. ORI_PIC
        comparison_panels.append(add_label_to_image(orig_img, "ORI_PIC"))
        
        # 5c. 加载 Ground Truth
        gt_color = np.zeros_like(orig_img) # 默认黑色
        if gt_path and gt_path.exists():
            gt_gray = cv2.imread(str(gt_path), cv2.IMREAD_GRAYSCALE)
            if gt_gray is not None:
                if gt_gray.shape != (h, w):
                    gt_gray = cv2.resize(gt_gray, (w, h), interpolation=cv2.INTER_NEAREST)
                gt_color = color_lut[gt_gray]
            else:
                logging.warning(f"\n无法读取GT图像 {gt_path}，使用黑色占位符。")
        else:
            logging.warning(f"\nGT图像未找到 (预期路径: {gt_path})，使用黑色占位符。")
        
        # 2. GT Panel (Mask or Overlay)
        if use_gt_overlay:
            gt_panel = cv2.addWeighted(orig_img, 0.6, gt_color, 0.4, 0)
        else:
            gt_panel = gt_color
        comparison_panels.append(add_label_to_image(gt_panel, gt_panel_label))

        # 5d. 加载每个模型的预测结果
        for model_dir in model_dirs:
            pred_path = model_dir / source_subdir_name / image_name
            pred_img = np.zeros_like(orig_img) # 默认黑色
            
            if pred_path.exists():
                img = cv2.imread(str(pred_path))
                if img is not None:
                    if img.shape != (h, w, 3):
                        logging.warning(f"\n预测图 {pred_path.name} 尺寸 {img.shape} 与原图 {(h,w,3)} 不匹配。正在调整大小...")
                        img = cv2.resize(img, (w, h), interpolation=cv2.INTER_AREA)
                    pred_img = img
                else:
                    logging.warning(f"\n无法读取预测图 {pred_path}，使用黑色占位符。")
            else:
                logging.warning(f"\n预测图未找到 {pred_path}，使用黑色占位符。")
                
            short_label = model_dir.name.split('_')[0]
            comparison_panels.append(add_label_to_image(pred_img, short_label))
        
        # 5e. 拼接并保存
        try:
            final_image = cv2.hconcat(comparison_panels)
            save_stem = Path(image_name).stem
            save_path = output_dir / (save_stem + ".png")
            cv2.imwrite(str(save_path), final_image)
                
        except Exception as e:
            logging.error(f"\n--- 拼接或保存 {image_name} 失败: {e} ---")
            logging.error(f"面板尺寸 (H, W, C): {[p.shape for p in comparison_panels]}")
            logging.warning("请检查所有图像是否具有相同的高度。")

    logging.info(f"--- ✅ '{output_dir_name}' 生成完毕 ---")


# --- [修改] 主函数 (重构) ---
def main():
    parser = argparse.ArgumentParser(
        description="将所有模型的预测结果 (Mask 和 Yolo) 与原图和GT图拼接成对比图。"
    )
    # 模式参数已移除
    parser.add_argument(
        "--pt_name",
        type=str,
        default="all",
        help="您想比较的权重文件名称 (例如 'best.pt' 或 'epoch100.pt'，'all'表示对全部权重进行处理)"
    )
    args = parser.parse_args()

    base_dir = config.PREDICT_BEST_MODEL_DIR

    # 1. 加载颜色LUT (只需一次)
    color_lut = load_color_lut(config.DATASET_YAML_PATH)
    if color_lut is None:
        logging.error("无法加载颜色LUT，脚本终止。")
        return

    # 2. 定义通用文件模式
    pt_name_raw = args.pt_name.replace('.pt', '')
    all_image_types = ["*.jpg", "*.jpeg", "*.png", "*.bmp", "*.tif", "*.tiff"]

    # 3. --- 执行 "Mask" 模式 ---
    mask_patterns = []
    if args.pt_name.lower() == "all":
        mask_patterns = ["*.png"]
    else:
        mask_patterns = [f"*_{pt_name_raw}.png"]
    
    generate_comparison_images(
        base_dir=base_dir,
        output_dir_name="compare_all_masks",
        source_subdir_name="predicted_color_masks",
        gt_panel_label="Ground_Truth",
        use_gt_overlay=False,
        file_glob_patterns=mask_patterns,
        color_lut=color_lut
    )

    # 4. --- 执行 "Yolo" 模式 ---
    yolo_patterns = []
    if args.pt_name.lower() == "all":
        yolo_patterns = all_image_types
    else:
        yolo_patterns = [
            f"*_{pt_name_raw}.jpg", f"*_{pt_name_raw}.jpeg", 
            f"*_{pt_name_raw}.png", f"*_{pt_name_raw}.bmp",
            f"*_{pt_name_raw}.tif", f"*_{pt_name_raw}.tiff"
        ]

    generate_comparison_images(
        base_dir=base_dir,
        output_dir_name="compare_all_masks_Yolo",
        source_subdir_name="prediction",
        gt_panel_label="GT_Overlay",
        use_gt_overlay=True,
        file_glob_patterns=yolo_patterns,
        color_lut=color_lut
    )

    logging.info("--- 🏁 所有对比图任务执行完毕 ---")


if __name__ == "__main__":
    main()