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

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

import os
import xarray as xr
from rioxarray import open_rasterio
import numpy as np


def sr2rgb(red_path: str, green_path: str, blue_path: str, output_path: str) -> None:
    """
    将红、绿、蓝三个单波段地表反射率图像合成为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}")

    # 读取三个波段数据
    red_band = open_rasterio(red_path, masked=True).squeeze(dim="band", drop=True)
    green_band = open_rasterio(green_path, masked=True).squeeze(dim="band", drop=True)
    blue_band = open_rasterio(blue_path, masked=True).squeeze(dim="band", drop=True)
    # 暂存元数据
    y_coords = red_band.y
    x_coords = red_band.x
    crs = red_band.rio.crs
    transform = red_band.rio.transform()

    def stretch_band(band):
        """
        线性拉伸到0-255范围
        """
        # 处理NaN值与负值
        band_no_nan = np.where(np.isnan(band), 0, band)
        band_no_nan = np.where(band_no_nan < 0, 0, band_no_nan)
        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.values)
    green_stretched = stretch_band(green_band.values)
    blue_stretched = stretch_band(blue_band.values)
    # 合并三个波段为RGB图像
    rgb_array = xr.DataArray(
        np.dstack((red_stretched, green_stretched, blue_stretched)),
        dims=("y", "x", "band"),
        coords={"band": [1, 2, 3], "y": y_coords, "x": x_coords},
    )
    # 转置维度顺序以符合rioxarray要求
    rgb_array = rgb_array.transpose("band", "y", "x")
    # 写入元数据
    rgb_array.rio.write_crs(crs, inplace=True)
    rgb_array.rio.write_transform(transform, inplace=True)
    rgb_array.rio.write_nodata(0, inplace=True)
    # 保存为TIFF文件
    rgb_array.rio.to_raster(output_path, dtype="uint8")
    print(f"RGB图像已保存到: {output_path}")
    return


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)