#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import cv2
import numpy as np
import os
import argparse
import matplotlib.pyplot as plt

def analyze_gray_sky_ratio(image_path, debug=False, save_result=False):
    """
    专门针对灰色天空图片的分析算法
    
    参数:
        image_path: 图片路径
        debug: 是否显示处理过程中的图像，用于调试
        save_result: 是否保存处理结果图像
        
    返回:
        sky_ratio: 天空占比（0-1之间的浮点数）
    """
    # 读取图片
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"无法读取图片: {image_path}")
    
    # 获取图片文件名（不带路径和扩展名）
    filename = os.path.splitext(os.path.basename(image_path))[0]
    
    # 转换为HSV色彩空间（更适合颜色分割）
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    
    # 提取图像的各个通道
    h, s, v = cv2.split(hsv)
    
    # 灰色天空特征：饱和度低（灰色是低饱和度的颜色）
    # 亮度适中到较高（相对于黑色路面）
    low_saturation_mask = s < 30  # 低饱和度阈值
    medium_high_value_mask = v > 120  # 中高亮度阈值
    
    # 组合掩码：低饱和度且中高亮度的区域可能是灰色天空
    sky_mask = np.logical_and(low_saturation_mask, medium_high_value_mask).astype(np.uint8) * 255
    
    # 应用形态学操作
    kernel = np.ones((5, 5), np.uint8)
    sky_mask = cv2.morphologyEx(sky_mask, cv2.MORPH_OPEN, kernel)
    sky_mask = cv2.morphologyEx(sky_mask, cv2.MORPH_CLOSE, kernel)
    
    # 假设天空通常位于图像上部
    height, width = sky_mask.shape
    sky_region_height = int(height * 0.6)  # 假设天空在上部60%的区域
    
    # 保留上部区域的掩码
    upper_mask = np.zeros_like(sky_mask)
    upper_mask[:sky_region_height, :] = 255
    
    # 结合上部区域限制
    sky_mask = cv2.bitwise_and(sky_mask, upper_mask)
    
    # 使用连通区域分析，去除小的噪点区域
    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(sky_mask, connectivity=8)
    
    # 过滤小的连通区域
    min_size = 1000  # 最小连通区域大小
    filtered_sky_mask = np.zeros_like(sky_mask)
    
    # 从索引1开始，因为0是背景
    for i in range(1, num_labels):
        if stats[i, cv2.CC_STAT_AREA] >= min_size:
            filtered_sky_mask[labels == i] = 255
    
    # 计算天空区域占比
    total_pixels = height * width
    sky_pixels = np.sum(filtered_sky_mask == 255)
    sky_ratio = sky_pixels / total_pixels
    
    # 在原图上标记天空区域
    result = img.copy()
    overlay = img.copy()
    overlay[filtered_sky_mask > 0] = [0, 255, 255]  # 用黄色标记天空区域
    cv2.addWeighted(overlay, 0.4, img, 0.6, 0, result)  # 半透明效果
    
    # 显示检测结果信息
    cv2.putText(result, f"Sky Ratio: {sky_ratio:.2%}", (10, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
    
    # 调试模式：显示处理过程图像
    if debug:
        plt.figure(figsize=(15, 10))
        
        plt.subplot(231)
        plt.title("Original Image")
        plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        
        plt.subplot(232)
        plt.title("Saturation Channel")
        plt.imshow(s, cmap='gray')
        
        plt.subplot(233)
        plt.title("Value Channel")
        plt.imshow(v, cmap='gray')
        
        plt.subplot(234)
        plt.title("Initial Sky Mask")
        plt.imshow(sky_mask, cmap='gray')
        
        plt.subplot(235)
        plt.title("Filtered Sky Mask")
        plt.imshow(filtered_sky_mask, cmap='gray')
        
        plt.subplot(236)
        plt.title("Sky Detection Result")
        plt.imshow(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
        
        plt.tight_layout()
        plt.show()
    
    # 保存结果
    if save_result:
        result_dir = "results"
        os.makedirs(result_dir, exist_ok=True)
        output_path = os.path.join(result_dir, f"{filename}_gray_sky_result.jpg")
        cv2.imwrite(output_path, result)
        print(f"结果已保存至: {output_path}")
    
    return sky_ratio

def main():
    parser = argparse.ArgumentParser(description='分析图片中灰色天空区域占比')
    # parser.add_argument('--image_path', default='./res/bar/image_20250525_085235.png', type=str, help='图片路径')
    parser.add_argument('--image_path', default='./res/bar/image_20250525_085240.png', type=str, help='图片路径')
    parser.add_argument('--debug', default=False, action='store_true', help='显示处理过程图像')
    parser.add_argument('--save', action='store_true', help='保存处理结果图像')
    args = parser.parse_args()
    
    try:
        sky_ratio = analyze_gray_sky_ratio(args.image_path, args.debug, args.save)
        print(f"灰色天空区域占比: {sky_ratio:.2%}")
    except Exception as e:
        print(f"错误: {e}")

if __name__ == "__main__":
    main()