import xdem
import rasterio
import numpy as np

def calc_cosine_i(solar_zn, solar_az, aspect, slope):
    """Generate cosine i image. The cosine of the incidence angle (i) is
       defined as the angle between the normal to the pixel surface
       and the solar zenith direction.
       All input geometry units must be in radians.

    Args:
        solar_az (numpy.ndarray): Solar azimuth angle.
        solar_zn (numpy.ndarray): Solar zenith angle.
        aspect (numpy.ndarray): Ground aspect.
        slope (numpy.ndarray): Ground slope.

    Returns:
        numpy.ndarray: Cosine i image.
    """
    relative_az = aspect - solar_az
    cosine_i = (np.cos(solar_zn) * np.cos(slope) +
                np.sin(solar_zn) * np.sin(slope) * np.cos(relative_az))
    return cosine_i

# ========== 1. 读取 DEM ==========
dem = xdem.DEM("E:/is2/yaopu/DEM/dsm.tif")

# ========== 2. 计算坡度、坡向 ==========
slope = xdem.terrain.slope(dem, method='ZevenbergThorne')
aspect = xdem.terrain.aspect(dem)

# ========== 3. 获取空间参考信息 ==========
transform = dem.transform
crs = dem.crs

# ========== 4. 转换为 numpy 数组 ==========
slope_data = slope.data
aspect_data = aspect.data

# ========== 5. 太阳角度（单位：度，根据实际飞行时间和地理位置设置） ==========
solar_zn_deg = 30.0   # 太阳天顶角（从天顶到太阳的夹角），例如 30°
solar_az_deg = 180.0  # 太阳方位角（从北顺时针），例如 180°（正南）

# 转换为弧度（因为三角函数需要弧度）
solar_zn_rad = np.deg2rad(solar_zn_deg)
solar_az_rad = np.deg2rad(solar_az_deg)

# 坡度和坡向也转换为弧度
slope_rad = np.deg2rad(slope_data)
aspect_rad = np.deg2rad(aspect_data)

# ========== 6. 计算 cosine_i ==========
cosine_i = calc_cosine_i(solar_zn_rad, solar_az_rad, aspect_rad, slope_rad)

# ========== 7. 处理无效值（NaN） ==========
# 注意：将 NaN 替换为 0 可能会与真实值为 0 的像元混淆（如水平面坡度为0，阴影边缘cosine_i=0）。
# 如需严格区分，建议使用特殊值（如 -9999）作为 NoData。
slope_data = np.nan_to_num(slope_data, nan=0.0)
aspect_data = np.nan_to_num(aspect_data, nan=0.0)
cosine_i = np.nan_to_num(cosine_i, nan=0.0)

# ========== 8. 堆叠为多波段数组（坡度、坡向、cosine_i） ==========
stacked = np.stack([slope_data, aspect_data, cosine_i], axis=0)

# ========== 9. 输出 ENVI 格式文件 ==========
output_path = "E:/is2/yaopu/DEM/slope_aspect_cosine.dat"

with rasterio.open(
    output_path,
    'w',
    driver='ENVI',
    height=stacked.shape[1],
    width=stacked.shape[2],
    count=stacked.shape[0],
    dtype=stacked.dtype,
    crs=crs,
    transform=transform,
    nodata=0.0,           # 与上面替换的 NaN 值一致
) as dst:
    dst.write(stacked)

print(f"已保存多波段 ENVI 文件：{output_path}")
print(f"波段顺序：1-坡度(°), 2-坡向(°), 3-cosine_i")