"""
将 COG 格式的 Red, Green, Blue 单波段地表反射率图像合成 RGB 图像

1. 对比度线性拉伸至 0-255;
2. 合并波段;
3. 保存为 uint8 格式 RGB 图像;
"""

import os
import rasterio as rio
import numpy as np

def sr2rgb(red_path, green_path, blue_path, output_path):
    """
    将红、绿、蓝三个单波段地表反射率图像合成为RGB图像
    
    参数:
        red_path (str): 红色波段文件路径
        green_path (str): 绿色波段文件路径  
        blue_path (str): 蓝色波段文件路径
        output_path (str): 输出RGB图像路径
    """
    # 检查文件是否存在
    for path in [red_path, green_path, blue_path]:
        if not os.path.exists(path):
            raise FileNotFoundError(f"文件不存在: {path}")
    
    print(f"正在处理文件:")
    print(f"  红波段: {red_path}")
    print(f"  绿波段: {green_path}")
    print(f"  蓝波段: {blue_path}")
    
    # 读取三个波段数据
    with rio.open(red_path) as red_src:
        red_band = red_src.read(1)
        profile = red_src.profile
        print(f"红波段形状: {red_band.shape}, 数据类型: {red_band.dtype}")
        
    with rio.open(green_path) as green_src:
        green_band = green_src.read(1)
        print(f"绿波段形状: {green_band.shape}, 数据类型: {green_band.dtype}")
        
    with rio.open(blue_path) as blue_src:
        blue_band = blue_src.read(1)
        print(f"蓝波段形状: {blue_band.shape}, 数据类型: {blue_band.dtype}")
    
    # 线性拉伸到0-255范围
    def stretch_band(band):
        # 处理NaN值
        band_no_nan = np.where(np.isnan(band), 0, band)
        band_min = np.min(band_no_nan)
        band_max = np.max(band_no_nan)
        
        # 避免除零错误
        if band_max == band_min:
            stretched = np.zeros_like(band_no_nan, dtype=np.uint8)
        else:
            stretched = ((band_no_nan - band_min) / (band_max - band_min) * 255).astype(np.uint8)
        return stretched
    
    red_stretched = stretch_band(red_band)
    green_stretched = stretch_band(green_band)
    blue_stretched = stretch_band(blue_band)
    
    # 合并三个波段为RGB图像
    rgb_array = np.dstack((red_stretched, green_stretched, blue_stretched))
    
    # 更新元数据
    profile.update(
        dtype=rio.uint8,
        count=3,
        photometric='RGB',
        nodata=0  # 设置nodata为0，因为uint8的范围是0-255
    )
    
    # 保存RGB图像
    with rio.open(output_path, 'w', **profile) as dst:
        for i in range(3):
            dst.write(rgb_array[:, :, i], i + 1)
    
    print(f"RGB图像已保存到: {output_path}")


if __name__ == "__main__":
    # tif_dir = "D:\\NASA_EarthData_Script\\data\\HLS\\2024\\2024012"
    # red_path = os.path.join(tif_dir, "HLS.S30.T49RGP.2024012T031101.v2.0.RED.subset.tif")
    # green_path = os.path.join(tif_dir, "HLS.S30.T49RGP.2024012T031101.v2.0.GREEN.subset.tif")
    # blue_path = os.path.join(tif_dir, "HLS.S30.T49RGP.2024012T031101.v2.0.BLUE.subset.tif")
    # output_path = os.path.join(tif_dir, "HLS.S30.T49RGP.2024012T031101.v2.0.RGB.tif")

    tif_dir = "D:\\NASA_EarthData_Script\\data\\HLS\\2025\\2025011"
    red_path = os.path.join(tif_dir, "HLS.S30.T49RGP.2025011T031009.v2.0.RED.subset.tif")
    green_path = os.path.join(tif_dir, "HLS.S30.T49RGP.2025011T031009.v2.0.GREEN.subset.tif")
    blue_path = os.path.join(tif_dir, "HLS.S30.T49RGP.2025011T031009.v2.0.BLUE.subset.tif")
    output_path = os.path.join(tif_dir, "HLS.S30.T49RGP.2025011T031009.v2.0.RGB.tif")
    
    sr2rgb(red_path, green_path, blue_path, output_path)