#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
DEM地形处理工具 - 计算坡度、坡向和太阳入射角余弦

该工具从DEM文件中提取坐标信息，根据给定时间计算太阳位置，
并生成坡度、坡向和cosine_i的多波段ENVI文件。

作者: BRDF_GUI Team
版本: 2.0.0
"""

import xdem
import rasterio
import numpy as np
import pandas as pd
import argparse
import sys
from datetime import datetime
from pathlib import Path
import pvlib.solarposition as solarposition
from rasterio.crs import CRS
import pyproj


def calc_cosine_i(solar_zn, solar_az, aspect, slope):
    """
    计算入射角余弦（cosine i）。
    
    入射角余弦定义为地表法线与太阳天顶方向之间夹角的余弦值。
    所有角度输入必须为弧度。

    Args:
        solar_az (numpy.ndarray): 太阳方位角（弧度）
        solar_zn (numpy.ndarray): 太阳天顶角（弧度）
        aspect (numpy.ndarray): 地面坡向（弧度）
        slope (numpy.ndarray): 地面坡度（弧度）

    Returns:
        numpy.ndarray: 入射角余弦图像
    """
    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


def get_dem_center_coords(dem):
    """
    获取DEM图像中心点的经纬度坐标。

    Args:
        dem: xdem.DEM对象

    Returns:
        tuple: (纬度, 经度) - 单位：度
    """
    # 获取DEM的范围边界
    bounds = dem.bounds  # (left, bottom, right, top)
    center_x = (bounds.left + bounds.right) / 2
    center_y = (bounds.bottom + bounds.top) / 2

    # 如果DEM是投影坐标系，需要转换为WGS84经纬度
    if dem.crs is not None and not dem.crs.is_geographic:
        # 创建从投影坐标系到WGS84的转换
        transformer = pyproj.Transformer.from_crs(dem.crs, CRS.from_epsg(4326), always_xy=True)
        lon, lat = transformer.transform(center_x, center_y)
    else:
        # 已经是地理坐标系（经纬度）
        lon, lat = center_x, center_y

    return lat, lon


def parse_datetime(datetime_str):
    """
    解析日期时间字符串。

    支持的格式：
    - YYYY-MM-DD HH:MM:SS
    - YYYY-MM-DDTHH:MM:SS
    - YYYYMMDD_HHMMSS

    Args:
        datetime_str (str): 日期时间字符串

    Returns:
        datetime: datetime对象（UTC时间）
    """
    formats = [
        "%Y-%m-%d %H:%M:%S",
        "%Y-%m-%dT%H:%M:%S",
        "%Y%m%d_%H%M%S",
        "%Y-%m-%d",
    ]

    for fmt in formats:
        try:
            dt = datetime.strptime(datetime_str, fmt)
            # 如果只输入了日期，默认时间为正午12:00
            if fmt == "%Y-%m-%d":
                dt = dt.replace(hour=12, minute=0, second=0)
            return dt
        except ValueError:
            continue

    raise ValueError(f"无法解析日期时间: {datetime_str}. 支持的格式: YYYY-MM-DD HH:MM:SS, YYYY-MM-DDTHH:MM:SS, YYYYMMDD_HHMMSS")


def get_user_input_datetime():
    """
    交互式获取用户输入的日期和时间。

    Returns:
        datetime: datetime对象（UTC时间）
    """
    print("请输入观测日期和时间（UTC时间）：")

    try:
        year = int(input("  年份 (YYYY): "))
        month = int(input("  月份 (MM): "))
        day = int(input("  日期 (DD): "))
        hour = int(input("  小时 (HH, 24小时制): "))
        minute = int(input("  分钟 (MM): "))
        second = int(input("  秒 (SS): "))

        dt = datetime(year, month, day, hour, minute, second)
        print(f"\n使用的时间: {dt.strftime('%Y-%m-%d %H:%M:%S')} UTC")
        return dt
    except ValueError as e:
        print(f"输入错误: {e}")
        print("将使用默认时间 2024-06-15 10:00:00 UTC")
        return datetime(2024, 6, 15, 10, 0, 0)


def calc_solar_angles(lat, lon, dt):
    """
    使用pvlib计算太阳天顶角和方位角。

    Args:
        lat (float): 纬度（度）
        lon (float): 经度（度）
        dt (datetime): datetime对象（UTC时间）

    Returns:
        tuple: (太阳天顶角, 太阳方位角) - 单位：度
               天顶角: 0=天顶, 90=地平线
               方位角: 从北顺时针，0=北，90=东，180=南，270=西
    """
    # 创建时间索引
    times = pd.DatetimeIndex([dt])

    # 使用pvlib计算太阳位置
    solpos = solarposition.get_solarposition(times, lat, lon)

    solar_zn_deg = solpos['zenith'].values[0]      # 天顶角
    solar_az_deg = solpos['azimuth'].values[0]    # 方位角

    return solar_zn_deg, solar_az_deg


def write_envi_output(output_path, data, transform, crs, metadata=None):
    """
    将数据写入ENVI格式文件。

    Args:
        output_path (str): 输出文件路径
        data (numpy.ndarray): 数据数组 (bands, rows, cols)
        transform: 地理变换参数
        crs: 坐标参考系
        metadata (dict, optional): 额外的元数据
    """
    with rasterio.open(
        output_path,
        'w',
        driver='ENVI',
        height=data.shape[1],
        width=data.shape[2],
        count=data.shape[0],
        dtype=data.dtype,
        crs=crs,
        transform=transform,
        nodata=0.0,
    ) as dst:
        dst.write(data)

        # 写入ENVI头文件元数据
        if metadata:
            dst.update_tags(**metadata)


def process_dem(dem_path, output_path, datetime_input=None, interactive=False):
    """
    主处理函数：读取DEM，计算坡度坡向，计算太阳角度，输出结果。

    Args:
        dem_path (str): DEM输入文件路径
        output_path (str): 输出文件路径
        datetime_input (str, optional): 日期时间字符串，如果为None则使用交互模式
        interactive (bool): 是否使用交互式时间输入

    Returns:
        dict: 处理结果信息
    """
    # ========== 1. 读取 DEM ==========
    print(f"读取DEM文件: {dem_path}")
    dem = xdem.DEM(dem_path)

    # ========== 2. 计算坡度、坡向 ==========
    print("计算坡度和坡向...")
    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. 获取DEM中心点坐标并计算太阳角度 ==========
    print("\n获取DEM坐标信息...")
    latitude, longitude = get_dem_center_coords(dem)
    print(f"  中心点纬度: {latitude:.6f}°")
    print(f"  中心点经度: {longitude:.6f}°")
    print(f"  坐标系: WGS84")

    # 获取观测时间
    if datetime_input:
        observation_time = parse_datetime(datetime_input)
        print(f"\n使用命令行提供的时间: {observation_time.strftime('%Y-%m-%d %H:%M:%S')} UTC")
    elif interactive or not sys.stdin.isatty():
        observation_time = get_user_input_datetime()
    else:
        # 默认时间（如果非交互式且未提供时间）
        observation_time = datetime(2024, 6, 15, 10, 0, 0)
        print(f"\n使用默认时间: {observation_time.strftime('%Y-%m-%d %H:%M:%S')} UTC")

    # 计算太阳角度
    print("\n计算太阳位置...")
    solar_zn_deg, solar_az_deg = calc_solar_angles(latitude, longitude, observation_time)

    print(f"  太阳天顶角: {solar_zn_deg:.2f}°")
    print(f"  太阳高度角: {90 - solar_zn_deg:.2f}°")
    print(f"  太阳方位角: {solar_az_deg:.2f}°")

    # 转换为弧度
    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 ==========
    print("\n计算入射角余弦 (cosine_i)...")
    cosine_i = calc_cosine_i(solar_zn_rad, solar_az_rad, aspect_rad, slope_rad)

    # ========== 7. 处理无效值（NaN） ==========
    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. 堆叠为多波段数组 ==========
    stacked = np.stack([slope_data, aspect_data, cosine_i], axis=0)

    # ========== 9. 准备元数据 ==========
    metadata = {
        'description': 'Slope, Aspect, and Cosine_i derived from DEM',
        'band_names': 'slope(deg),aspect(deg),cosine_i',
        'solar_zenith_angle': str(solar_zn_deg),
        'solar_azimuth_angle': str(solar_az_deg),
        'solar_elevation_angle': str(90 - solar_zn_deg),
        'observation_time_utc': observation_time.strftime('%Y-%m-%d %H:%M:%S'),
        'dem_center_lat': str(latitude),
        'dem_center_lon': str(longitude),
    }

    # ========== 10. 输出 ENVI 格式文件 ==========
    print(f"\n保存结果到: {output_path}")
    write_envi_output(output_path, stacked, transform, crs, metadata)

    print("\n处理完成!")
    print(f"  输出文件: {output_path}")
    print(f"  波段顺序: 1-坡度(°), 2-坡向(°), 3-cosine_i")

    return {
        'output_path': output_path,
        'solar_zenith': solar_zn_deg,
        'solar_azimuth': solar_az_deg,
        'center_lat': latitude,
        'center_lon': longitude,
        'observation_time': observation_time,
    }


def main():
    """主函数 - 命令行入口"""
    parser = argparse.ArgumentParser(
        description='DEM地形处理工具 - 计算坡度、坡向和太阳入射角余弦',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
使用示例:
  # 基本用法（交互式输入时间）
  python slope_aspectV2.py -i dem.tif -o output.dat

  # 通过命令行指定时间
  python slope_aspectV2.py -i dem.tif -o output.dat -t "2024-06-15 10:30:00"

  # 仅指定日期（默认使用正午12:00）
  python slope_aspectV2.py -i dem.tif -o output.dat -t "2024-06-15"

  # 显示详细处理信息
  python slope_aspectV2.py -i dem.tif -o output.dat -t "2024-06-15 10:30:00" -v

支持的日期时间格式:
  - "YYYY-MM-DD HH:MM:SS" (推荐)
  - "YYYY-MM-DDTHH:MM:SS" (ISO格式)
  - "YYYYMMDD_HHMMSS"
  - "YYYY-MM-DD" (默认使用12:00:00)
        '''
    )

    parser.add_argument('-i', '--input', required=True,
                        help='输入DEM文件路径 (支持GeoTIFF等格式)')

    parser.add_argument('-o', '--output', required=True,
                        help='输出ENVI文件路径 (.dat)')

    parser.add_argument('-t', '--time',
                        help='观测日期时间 (UTC)，格式: "YYYY-MM-DD HH:MM:SS"。'
                             '如果不提供，将进入交互式输入模式')

    parser.add_argument('--interactive', action='store_true',
                        help='强制使用交互式时间输入模式')

    parser.add_argument('-v', '--verbose', action='store_true',
                        help='显示详细处理信息')

    parser.add_argument('--version', action='version', version='%(prog)s 2.0.0')

    args = parser.parse_args()

    # 验证输入文件存在
    if not Path(args.input).exists():
        print(f"错误: 输入文件不存在: {args.input}", file=sys.stderr)
        sys.exit(1)

    # 创建输出目录（如果不存在）
    output_dir = Path(args.output).parent
    if not output_dir.exists():
        output_dir.mkdir(parents=True, exist_ok=True)
        print(f"创建输出目录: {output_dir}")

    try:
        # 执行处理
        result = process_dem(
            dem_path=args.input,
            output_path=args.output,
            datetime_input=args.time,
            interactive=args.interactive
        )

        if args.verbose:
            print("\n详细结果:")
            for key, value in result.items():
                print(f"  {key}: {value}")

        sys.exit(0)

    except Exception as e:
        print(f"\n处理失败: {e}", file=sys.stderr)
        if args.verbose:
            import traceback
            traceback.print_exc()
        sys.exit(1)


if __name__ == '__main__':
    main()