# slope_aspect_img13
import numpy as np
import matplotlib.pyplot as plt
from scipy.ndimage import sobel
import matplotlib as mpl
import rasterio
from osgeo import gdal
from matplotlib.patches import FancyArrowPatch
from mpl_toolkits.mplot3d import proj3d
import os

# 自定义3D箭头类
class Arrow3D(FancyArrowPatch):
    def __init__(self, xs, ys, zs, *args, **kwargs):
        super().__init__((0,0), (0,0), *args, **kwargs)
        self._verts3d = xs, ys, zs

    def do_3d_projection(self, renderer=None):
        xs3d, ys3d, zs3d = self._verts3d
        xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, self.axes.M)
        self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
        return min(zs)

def read_dem_rasterio(filepath):
    """使用rasterio读取DEM数据"""
    try:
        with rasterio.open(filepath) as src:
            dem_data = src.read(1)
            transform = src.transform
            bounds = src.bounds
            crs = src.crs
            
            print(f"DEM信息:")
            print(f"  尺寸: {dem_data.shape}")
            print(f"  范围: {bounds}")
            print(f"  坐标系: {crs}")
            print(f"  高程范围: {np.nanmin(dem_data):.2f} - {np.nanmax(dem_data):.2f}")
            
            if np.isnan(dem_data).any():
                print("  检测到NaN值")
                dem_data = np.nan_to_num(dem_data, nan=np.nanmean(dem_data))
            
            rows, cols = dem_data.shape
            x = np.linspace(bounds.left, bounds.right, cols)
            y = np.linspace(bounds.bottom, bounds.top, rows)
            X, Y = np.meshgrid(x, y)
            
            return X, Y, dem_data
            
    except Exception as e:
        print(f"使用rasterio读取失败: {e}")
        return None
    
def read_slope_aspect_by_dem(dem_path, overall_3d_output_path=None) :
    # ---------------------- 核心：解决中文乱码配置 ----------------------
    setup_chinese_font()

    # 尝试读取DEM数据
    print("正在读取DEM文件...")
    X, Y, Z = read_dem_rasterio(dem_path)

    if X is None:
        raise FileNotFoundError(f"无法读取DEM文件: {dem_path}")

    # 检查数据有效性
    if Z.size == 0 or np.all(Z == 0):
        raise ValueError("DEM数据无效或全部为0")

    # 重采样
    if Z.shape[0] > 200 or Z.shape[1] > 200:
        print(f"原始DEM尺寸较大 ({Z.shape})，进行重采样...")
        from scipy.ndimage import zoom
        scale_factor = min(200/Z.shape[0], 200/Z.shape[1])
        Z = zoom(Z, scale_factor, order=1)
        X = zoom(X, scale_factor, order=1)
        Y = zoom(Y, scale_factor, order=1)
        print(f"重采样后尺寸: {Z.shape}")

    # ---------------------- 2. 计算坡度和坡向 ----------------------
    print("计算坡度和坡向...")

    dx_pixel = np.abs(X[0,1] - X[0,0]) * 111000
    dy_pixel = np.abs(Y[1,0] - Y[0,0]) * 111000

    dx = sobel(Z, axis=1) / (2 * dx_pixel)
    dy = sobel(Z, axis=0) / (2 * dy_pixel)

    slope_rad = np.arctan(np.sqrt(dx**2 + dy**2))
    slope_deg = slope_rad * (180 / np.pi)

    aspect_rad = np.arctan2(-dx, dy)
    aspect_deg = aspect_rad * (180 / np.pi)
    aspect_deg[aspect_deg < 0] += 360

    print(f"坡度范围: {np.min(slope_deg):.2f}° - {np.max(slope_deg):.2f}°")
    print(f"坡向范围: {np.min(aspect_deg):.2f}° - {np.max(aspect_deg):.2f}°")

    # ---------------------- 3. 计算整体坡向 ----------------------
    print("\n计算整体坡向...")

    def calculate_overall_aspect(aspect_deg, slope_deg, method='weighted_mean'):
        """计算整体坡向"""
        aspect_rad = np.deg2rad(aspect_deg)
        
        if method == 'weighted_mean':
            # 坡度加权平均法
            u = np.sin(aspect_rad)
            v = np.cos(aspect_rad)
            
            weights = slope_deg.flatten()
            weighted_u = np.nansum(u.flatten() * weights) / np.nansum(weights)
            weighted_v = np.nansum(v.flatten() * weights) / np.nansum(weights)
            
            weighted_aspect_rad = np.arctan2(weighted_u, weighted_v)
            weighted_aspect_deg = np.rad2deg(weighted_aspect_rad)
            
            if weighted_aspect_deg < 0:
                weighted_aspect_deg += 360
                
            weighted_strength = np.sqrt(weighted_u**2 + weighted_v**2)
            return weighted_aspect_deg, weighted_strength, "坡度加权平均法"
        
        elif method == 'vector_mean':
            # 向量平均法
            u = np.sin(aspect_rad)
            v = np.cos(aspect_rad)
            
            mean_u = np.nanmean(u)
            mean_v = np.nanmean(v)
            
            mean_aspect_rad = np.arctan2(mean_u, mean_v)
            mean_aspect_deg = np.rad2deg(mean_aspect_rad)
            
            if mean_aspect_deg < 0:
                mean_aspect_deg += 360
                
            vector_strength = np.sqrt(mean_u**2 + mean_v**2)
            return mean_aspect_deg, vector_strength, "向量平均法"

    # 使用坡度加权平均法计算整体坡向
    overall_aspect, overall_strength, method_name = calculate_overall_aspect(aspect_deg, slope_deg, 'weighted_mean')

    # 将整体坡向转换为方向描述
    if overall_aspect < 22.5 or overall_aspect >= 337.5:
        overall_direction = "北"
    elif 22.5 <= overall_aspect < 67.5:
        overall_direction = "东北"
    elif 67.5 <= overall_aspect < 112.5:
        overall_direction = "东"
    elif 112.5 <= overall_aspect < 157.5:
        overall_direction = "东南"
    elif 157.5 <= overall_aspect < 202.5:
        overall_direction = "南"
    elif 202.5 <= overall_aspect < 247.5:
        overall_direction = "西南"
    elif 247.5 <= overall_aspect < 292.5:
        overall_direction = "西"
    else:
        overall_direction = "西北"

    print(f"整体坡向 ({method_name}):")
    print(f"  角度: {overall_aspect:.1f}°")
    print(f"  方向: {overall_direction}")
    # print(f"  一致性: {overall_strength:.3f}")

    # ---------------------- 5. 3D可视化（俯视图，包含整体坡向和关键点坡向） ----------------------
    print("\n生成3D俯视图可视化...")
    fig3d, ax3d = plt.subplots(figsize=(16, 12), subplot_kw={"projection": "3d"})

    # 绘制地形曲面 - 俯视图需要更清晰的地形表现
    norm = mpl.colors.Normalize(vmin=np.percentile(slope_deg, 5), 
                            vmax=np.percentile(slope_deg, 95))

    plot_skip = max(2, Z.shape[0] // 60)  # 增加采样密度，使俯视图更清晰
    X_plot = X[::plot_skip, ::plot_skip]
    Y_plot = Y[::plot_skip, ::plot_skip]
    Z_plot = Z[::plot_skip, ::plot_skip]
    slope_plot = slope_deg[::plot_skip, ::plot_skip]

    surf = ax3d.plot_surface(
        X_plot, Y_plot, Z_plot,
        cmap="viridis_r",
        alpha=0.85,  # 增加透明度，使箭头更明显
        linewidth=0.1,  # 很细的网格线
        facecolors=plt.cm.viridis_r(norm(slope_plot)),
        zorder=1
    )

    # ---------------------- 绘制整体坡向箭头（中心位置，红色粗箭头） ----------------------
    center_x = (X.min() + X.max()) / 2
    center_y = (Y.min() + Y.max()) / 2

    # 整体坡向箭头长度
    arrow_length_overall = 0.15 * min(X.max()-X.min(), Y.max()-Y.min())

    # 计算整体坡向箭头方向
    scale_factor = 1.5
    overall_aspect_rad = np.deg2rad(overall_aspect)
    dx_overall = np.sin(overall_aspect_rad) * arrow_length_overall * scale_factor
    dy_overall = np.cos(overall_aspect_rad) * arrow_length_overall * scale_factor

    # 方案2：如果希望箭头在地形上方一定高度
    terrain_max_z = Z.max()  # 地形最高点
    float_height = 0.2 * terrain_max_z  # 在地形最高点上方20%的高度

    # 绘制整体坡向箭头（红色，粗）
    arrow_overall = Arrow3D(
        [center_x, center_x + dx_overall], 
        [center_y, center_y + dy_overall], 
        [terrain_max_z + float_height, terrain_max_z + float_height],  # 在地形上方
        mutation_scale=25,  # 稍大一些的箭头
        lw=5,  # 更粗的线
        arrowstyle='-|>',
        color='red',
        alpha=0.98,  # 更高的透明度
        zorder=20  # 更高的绘制顺序
    )
    ax3d.add_artist(arrow_overall)

    # 在整体坡向箭头起点添加标记
    ax3d.scatter(center_x, center_y, terrain_max_z + float_height, 
                color='red', s=120, edgecolor='white', linewidth=2, zorder=21)

    # ---------------------- 整体坡向面板放置在左上角 ----------------------
    # 计算左上角位置
    panel_x = X.min() + 0.02 * (X.max() - X.min())  # 左边留2%的边距
    panel_y = Y.max() - 0.02 * (Y.max() - Y.min())  # 上边留2%的边距
    panel_z = Z.max() + 0.5 * (Z.max() - Z.min())   # 进一步提高Z坐标，确保在视野内

    # 整体坡向面板内容
    overall_info = (
        f"整体坡向分析\n"
        f"角度: {overall_aspect:.1f}°\n"
        f"方向: {overall_direction}"
        # f"一致性: {overall_strength:.3f}"
    )

    # 绘制整体坡向面板（左上角）
    ax3d.text(panel_x, panel_y, panel_z,
            overall_info,
            fontsize=12, fontweight='bold',
            bbox=dict(facecolor='white', alpha=0.5, boxstyle="round,pad=0.5",
                    edgecolor='red', linewidth=2.5),
            ha='left', va='top', zorder=30)

    # ---------------------- 设置3D图参数 ----------------------
    ax3d.set_xlabel("经度 (X)", fontsize=12)
    ax3d.set_ylabel("纬度 (Y)", fontsize=12)
    ax3d.set_zlabel("高程 (m)", fontsize=12)

    # 获取文件名用于标题
    filename = os.path.basename(dem_path)
    ax3d.set_title(f"DEM三维俯视图 - 坡向分析\n文件: {filename}", 
                fontsize=14, fontweight='bold', pad=20)

    # 添加颜色条
    cbar = plt.colorbar(
        mpl.cm.ScalarMappable(norm=norm, cmap='viridis_r'),
        ax=ax3d, shrink=0.6, aspect=25, pad=0.15
    )
    cbar.set_label("坡度 (°)", fontsize=12)

    # ---------------------- 设置为俯视图（高仰角） ----------------------
    view_elev = 85  # 接近90度的仰角，俯视效果
    view_azim = overall_aspect + 180  # 从坡向的相反方向观看，可以看到坡面

    # 确保方位角在0-360度范围内
    view_azim = view_azim % 360

    print(f"\n设置3D俯视图视角:")
    print(f"  整体坡向: {overall_aspect:.1f}° ({overall_direction})")
    print(f"  视角方位角: {view_azim:.1f}°")
    print(f"  视角仰角: {view_elev:.1f}°")

    # 应用俯视图视角设置
    ax3d.view_init(elev=view_elev, azim=view_azim)

    # 调整相机距离，使视角更广
    ax3d.dist = 9.0  # 增加相机距离

    # 设置坐标轴范围，确保所有元素都显示
    z_min, z_max = Z.min(), Z.max()
    z_padding = 0.6 * (z_max - z_min)  # 适当增加Z轴范围
    ax3d.set_zlim(z_min - 0.1*z_padding, z_max + z_padding)

    # 设置XY轴范围
    x_margin = 0.1 * (X.max() - X.min())
    y_margin = 0.1 * (Y.max() - Y.min())
    ax3d.set_xlim(X.min() - x_margin, X.max() + x_margin)
    ax3d.set_ylim(Y.min() - y_margin, Y.max() + y_margin)

    plt.tight_layout()

    # 保存图片
    if not overall_3d_output_path : 
        overall_3d_output_path = dem_path.replace('.tif', '_slopeAspect_3D_overlook.png')
    plt.savefig(overall_3d_output_path, dpi=250, bbox_inches='tight', facecolor='white')
    # plt.show()

    os.remove(dem_path)
    print(f"\n3D俯视图已保存: {overall_3d_output_path}")
    print("分析完成！")
    print("="*60)

    return overall_3d_output_path

# ---------------------- 跨平台中文字体配置 ----------------------
def setup_chinese_font():
    """设置中文字体支持，兼容Windows、Linux、macOS"""
    import platform
    
    plt.rcParams['axes.unicode_minus'] = False  # 正确显示负号
    
    system = platform.system()
    
    # 尝试添加系统中文字体路径
    font_paths = []
    
    if system == 'Windows':
        # Windows 字体路径
        font_paths = [
            'C:/Windows/Fonts/simhei.ttf',      # 黑体
            'C:/Windows/Fonts/simkai.ttf',      # 楷体
            'C:/Windows/Fonts/simsun.ttc',      # 宋体
            'C:/Windows/Fonts/microsoftyahei.ttf',  # 微软雅黑
        ]
    elif system == 'Darwin':  # macOS
        # macOS 字体路径
        font_paths = [
            '/System/Library/Fonts/PingFang.ttc',    # 苹方
            '/System/Library/Fonts/STHeiti Light.ttc', # 华文黑体
            '/System/Library/Fonts/STHeiti Medium.ttc',
            '/Library/Fonts/Arial Unicode.ttf',      # Arial Unicode
        ]
    elif system == 'Linux':
        # Linux 字体路径
        font_paths = [
            '/usr/share/fonts/truetype/wqy/wqy-microhei.ttc',  # 文泉驿微米黑
            '/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc',    # 文泉驿正黑
            '/usr/share/fonts/truetype/arphic/uming.ttc',      # AR PL UMing
            '/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc',  # Noto
        ]
    
    # 尝试找到并添加第一个可用的中文字体
    font_added = False
    for font_path in font_paths:
        try:
            if os.path.exists(font_path):
                font_prop = mpl.font_manager.FontProperties(fname=font_path)
                font_name = font_prop.get_name()
                mpl.font_manager.fontManager.addfont(font_path)
                plt.rcParams['font.sans-serif'] = [font_name] + plt.rcParams['font.sans-serif']
                font_added = True
                print(f"已添加中文字体: {font_name} ({font_path})")
                break
        except Exception as e:
            print(f"添加字体 {font_path} 失败: {e}")
            continue
    
    # 如果以上方法都失败，使用通用备选方案
    if not font_added:
        if system == 'Windows':
            fallback_fonts = ['SimHei', 'Microsoft YaHei', 'KaiTi', 'FangSong']
        elif system == 'Darwin':
            fallback_fonts = ['PingFang SC', 'Hiragino Sans GB', 'Apple LiGothic Medium']
        else:  # Linux and others
            fallback_fonts = ['DejaVu Sans', 'WenQuanYi Micro Hei', 
                            'Noto Sans CJK SC', 'Heiti TC', 'AR PL UMing CN']
        
        current_fonts = plt.rcParams.get('font.sans-serif', [])
        plt.rcParams['font.sans-serif'] = fallback_fonts + current_fonts
        print(f"使用备选字体方案: {fallback_fonts[:2]}...")
    
    # 设置字体族
    plt.rcParams['font.family'] = 'sans-serif'

  
if __name__ == "__main__":
    dem_path = r'D:/devForBdzlWork/ai_project_v1/b3dm/o_dem_f1cb6f69_slopeAspect.tif'
    read_slope_aspect_by_dem(dem_path)