import rasterio
from rasterio.warp import reproject, Resampling
import numpy as np

def merge_angles_with_slope_aspect(angle_file, slope_aspect_file, output_file,
                                   resampling=Resampling.bilinear):
    """
    将坡度坡向波段添加到角度文件波段之后，输出合并的ENVI DAT文件。
    输出文件的空间范围、分辨率、投影与角度文件完全一致。
    掩膜基于角度文件的忽略值生成：只有所有角度波段均为忽略值(0.0)的像素才被设为忽略值。
    输出文件的忽略值统一设为0.0。

    Parameters
    ----------
    angle_file : str
        包含角度信息的ENVI DAT文件（如传感器角度、太阳角度等，至少1个波段）。
        该文件定义了输出的目标空间参数和掩膜基准。
    slope_aspect_file : str
        包含坡度坡向的ENVI DAT文件（通常2个波段）。
    output_file : str
        输出合并后的ENVI DAT文件路径（将自动生成.hdr头文件）。
    resampling : rasterio.warp.Resampling, optional
        重采样方法，默认为双线性（适用于连续参数）。
    """
    # 1. 从角度文件获取目标空间参数、数据、掩膜
    with rasterio.open(angle_file) as src_angle:
        # 目标空间参数
        dst_transform = src_angle.transform
        dst_width = src_angle.width
        dst_height = src_angle.height
        dst_crs = src_angle.crs

        # 读取角度文件所有波段
        angle_data = src_angle.read()  # shape: (bands, rows, cols)
        angle_count = src_angle.count

        # 获取角度文件的忽略值（如果未定义，默认为0.0，但应给出警告）
        angle_nodata = src_angle.nodata
        if angle_nodata is None:
            print("警告：角度文件未定义忽略值，将使用0.0作为默认忽略值。")
            angle_nodata = 0.0

        # 构建掩膜：所有波段均不等于忽略值视为有效
        mask = np.all(angle_data != angle_nodata, axis=0)

    # 2. 处理坡度坡向文件：重采样到目标空间
    with rasterio.open(slope_aspect_file) as src_sa:
        src_sa_data = src_sa.read()          # shape: (bands, src_sa_height, src_sa_width)
        src_sa_count = src_sa.count
        src_sa_crs = src_sa.crs
        src_sa_transform = src_sa.transform

        sa_resampled = np.empty((src_sa_count, dst_height, dst_width),
                                dtype=src_sa_data.dtype)

        reproject(
            source=src_sa_data,
            destination=sa_resampled,
            src_transform=src_sa_transform,
            src_crs=src_sa_crs,
            dst_transform=dst_transform,
            dst_crs=dst_crs,
            resampling=resampling
        )

    # 3. 合并波段：角度在前，坡度坡向在后
    merged_data = np.concatenate([angle_data, sa_resampled], axis=0)
    merged_count = merged_data.shape[0]

    # 4. 准备输出元数据（基于角度文件的profile）
    with rasterio.open(angle_file) as src_angle:
        out_profile = src_angle.profile.copy()
        output_nodata = 0.0   # 强制输出忽略值为0.0（用户要求）
        out_profile.update({
            'driver': 'ENVI',
            'height': dst_height,
            'width': dst_width,
            'transform': dst_transform,
            'crs': dst_crs,
            'count': merged_count,
            'dtype': merged_data.dtype,
            'nodata': output_nodata
        })

    # 5. 应用掩膜：将无效区域的像素设为忽略值0.0
    for b in range(merged_count):
        merged_data[b, ~mask] = output_nodata

    # 6. 写入输出文件
    with rasterio.open(output_file, 'w', **out_profile) as dst:
        dst.write(merged_data)

    print(f"处理完成！合并后的文件已保存至：{output_file}")
    print(f"总波段数：{merged_count}（角度 {angle_count} 波段 + 坡度坡向 {src_sa_count} 波段）")
    print(f"有效像素区域基于角度文件所有波段判断，忽略值统一设为 0.0。")

# 示例用法
if __name__ == "__main__":
    angle_file = r"E:\is2\yaopu\output\angel.dat"               # 角度DAT文件
    slope_aspect_file = r"E:\is2\yaopu\DEM\slope_aspect_cosine.dat"   # 原始坡度坡向DAT
    out_file = r"E:\is2\yaopu\DEM\angles_with_slope_aspect.dat" # 合并输出

    merge_angles_with_slope_aspect(angle_file, slope_aspect_file, out_file)