#!/usr/bin/env python3
"""
3D可视化脚本 - 显示经纬度高程数据

读取CSV文件中的经度、纬度、融合高程数据，
生成可交互的三维点图，输出为HTML格式。

数据格式要求：
经度 纬度 高程 融合高程

经纬度坐标需要除以10^-7进行单位转换。
"""

import pandas as pd
import plotly.graph_objects as go
import plotly.io as pio
import argparse
import sys
import numpy as np
from pathlib import Path

def load_and_process_data(file_path, lon_col=0, lat_col=1, alt_col=3, delimiter=None, sheet_name=0):
    """
    加载并处理CSV或Excel数据

    Args:
        file_path: 文件路径 (支持.csv, .xlsx, .xls)
        lon_col: 经度列索引 (默认0)
        lat_col: 纬度列索引 (默认1)
        alt_col: 高程列索引 (默认3，为融合高程)
        delimiter: CSV分隔符，None表示自动检测
        sheet_name: Excel工作表名称或索引 (默认0)

    Returns:
        pandas.DataFrame: 处理后的数据
    """
    try:
        file_path_obj = Path(file_path)

        # 根据文件扩展名选择读取方式
        if file_path_obj.suffix.lower() in ['.xlsx', '.xls']:
            # 读取Excel文件
            print(f"检测到Excel文件: {file_path}")
            df = pd.read_excel(file_path, sheet_name=sheet_name, header=None)
        else:
            # 读取CSV文件
            print(f"检测到CSV文件: {file_path}")
            df = pd.read_csv(file_path, delimiter=delimiter, header=None)

        # 验证数据列数
        if df.shape[1] < 4:
            raise ValueError(f"数据文件至少需要4列，经度、纬度、高程、融合高程。当前只有{df.shape[1]}列。")

        # 显示前几行数据以便调试
        print(f"数据预览 (前3行):")
        print(df.head(3))
        print()

        # 提取所需列
        lon_raw = df.iloc[:, lon_col]  # 经度（原始格式）
        lat_raw = df.iloc[:, lat_col]  # 纬度（原始格式）
        altitude = df.iloc[:, alt_col]  # 融合高程

        # 数据类型转换和清理
        try:
            lon_raw = pd.to_numeric(lon_raw, errors='coerce')
            lat_raw = pd.to_numeric(lat_raw, errors='coerce')
            altitude = pd.to_numeric(altitude, errors='coerce')
        except Exception as e:
            print(f"数据类型转换失败: {e}")
            print("尝试手动转换...")

            # 手动转换函数
            def safe_convert(value):
                try:
                    return float(value)
                except (ValueError, TypeError):
                    return np.nan

            lon_raw = lon_raw.apply(safe_convert)
            lat_raw = lat_raw.apply(safe_convert)
            altitude = altitude.apply(safe_convert)

        # 检查转换结果
        nan_count_lon = lon_raw.isna().sum()
        nan_count_lat = lat_raw.isna().sum()
        nan_count_alt = altitude.isna().sum()

        if nan_count_lon > 0 or nan_count_lat > 0 or nan_count_alt > 0:
            print(f"警告：发现无效数据点 - 经度: {nan_count_lon}, 纬度: {nan_count_lat}, 高程: {nan_count_alt}")
            print("无效数据点将被移除")

        # 移除NaN值
        valid_mask = ~(lon_raw.isna() | lat_raw.isna() | altitude.isna())
        lon_raw = lon_raw[valid_mask]
        lat_raw = lat_raw[valid_mask]
        altitude = altitude[valid_mask]

        if len(lon_raw) == 0:
            raise ValueError("没有有效的坐标数据")

        # 经纬度单位转换 (除以10^7)
        longitude = lon_raw / 1e7
        latitude = lat_raw / 1e7

        # 创建处理后的DataFrame
        processed_df = pd.DataFrame({
            'longitude': longitude,
            'latitude': latitude,
            'altitude': altitude
        })

        print(f"成功加载数据：{len(processed_df)} 个数据点")
        print(f"经度范围: {longitude.min():.6f} - {longitude.max():.6f}")
        print(f"纬度范围: {latitude.min():.6f} - {latitude.max():.6f}")
        print(f"高程范围: {altitude.min():.2f} - {altitude.max():.2f} 米")

        return processed_df

    except Exception as e:
        print(f"数据加载失败: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

def create_3d_plot(df, title="3D轨迹可视化", output_file="3d_view.html"):
    """
    创建三维散点图

    Args:
        df: 包含longitude, latitude, altitude列的DataFrame
        title: 图表标题
        output_file: 输出HTML文件路径
    """

    # 创建3D散点图
    fig = go.Figure(data=[go.Scatter3d(
        x=df['longitude'],
        y=df['latitude'],
        z=df['altitude'],
        mode='markers',
        marker=dict(
            size=3,
            color=df['altitude'],  # 颜色映射到高程
            colorscale='Viridis',  # 颜色方案
            colorbar=dict(title="高程 (m)"),
            opacity=0.8
        ),
        hovertemplate=(
            '经度: %{x:.6f}<br>' +
            '纬度: %{y:.6f}<br>' +
            '高程: %{z:.2f} m<br>' +
            '<extra></extra>'
        )
    )])

    # 设置布局
    fig.update_layout(
        title=title,
        scene=dict(
            xaxis_title='经度',
            yaxis_title='纬度',
            zaxis_title='高程 (m)',
            xaxis=dict(showgrid=True, gridcolor='lightgray'),
            yaxis=dict(showgrid=True, gridcolor='lightgray'),
            zaxis=dict(showgrid=True, gridcolor='lightgray'),
            # 设置视角
            camera=dict(
                eye=dict(x=1.5, y=1.5, z=1.5)  # 初始视角位置
            )
        ),
        margin=dict(l=0, r=0, b=0, t=40),
        paper_bgcolor='white',
        plot_bgcolor='white'
    )

    # 添加交互说明
    fig.add_annotation(
        text="鼠标拖拽旋转视角，滚轮缩放，双击重置",
        xref="paper", yref="paper",
        x=0.02, y=0.98,
        showarrow=False,
        bgcolor="rgba(255,255,255,0.8)",
        bordercolor="black",
        borderwidth=1
    )

    # 保存为HTML文件
    try:
        pio.write_html(fig, file=output_file, auto_open=False)
        print(f"3D可视化已保存到: {output_file}")
        print("在浏览器中打开该文件即可查看可交互的3D图表")
    except Exception as e:
        print(f"保存HTML文件失败: {e}")
        sys.exit(1)

    return fig

def main(input_file, output_file='3d_view.html', lon_col=2, lat_col=3, alt_col=5,
         delimiter=None, title='3D轨迹可视化', sheet_name=0):
    """
    直接调用函数生成3D可视化

    Args:
        input_file: 输入文件路径 (支持.csv, .xlsx, .xls)
        output_file: 输出HTML文件路径
        lon_col: 经度列索引 (默认: 0)
        lat_col: 纬度列索引 (默认: 1)
        alt_col: 融合高程列索引 (默认: 3)
        delimiter: CSV分隔符 (默认: 自动检测)
        title: 图表标题 (默认: '3D轨迹可视化')
        sheet_name: Excel工作表名称或索引 (默认: 0)
    """

    # 检查输入文件
    input_path = Path(input_file)
    if not input_path.exists():
        print(f"错误：输入文件不存在: {input_path}")
        return False

    print(f"正在处理文件: {input_path}")
    print(f"输出文件: {output_file}")
    print("-" * 50)

    try:
        # 加载和处理数据
        df = load_and_process_data(
            file_path=input_path,
            lon_col=lon_col,
            lat_col=lat_col,
            alt_col=alt_col,
            delimiter=delimiter,
            sheet_name=sheet_name
        )

        # 创建3D图表
        fig = create_3d_plot(
            df=df,
            title=title,
            output_file=output_file
        )

        print("-" * 50)
        print("处理完成！")
        print(f"\n✅ 3D可视化已保存到: {output_file}")
        print("\n使用说明:")
        print("- 在浏览器中打开生成的HTML文件")
        print("- 鼠标左键拖拽：旋转视角")
        print("- 鼠标滚轮：缩放视图")
        print("- 双击：重置到初始视角")
        print("- 右键拖拽：平移视图")

        return True

    except Exception as e:
        print(f"❌ 处理失败: {e}")
        import traceback
        traceback.print_exc()
        return False

# 命令行接口（保持兼容性）
def cli_main():
    """命令行接口"""
    parser = argparse.ArgumentParser(
        description="3D轨迹可视化工具 - 将CSV数据转换为可交互的3D点图",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # CSV文件
  python 3Dview.py data.csv
  python 3Dview.py data.csv -o custom_output.html
  python 3Dview.py data.csv --lon-col 1 --lat-col 2 --alt-col 4

  # Excel文件
  python 3Dview.py trajectory.xlsx
  python 3Dview.py data.xlsx --sheet-name "Sheet1"

数据格式:
  文件应包含至少4列：经度 纬度 高程 融合高程
  支持格式: CSV (.csv), Excel (.xlsx, .xls)
  经纬度将自动除以10^7进行单位转换

直接调用示例:
  from 3Dview import main
  main('data.xlsx', sheet_name='Sheet1')  # Excel文件
  main('data.csv', lon_col=0, lat_col=1, alt_col=3)  # CSV文件
        """
    )

    parser.add_argument('input_file', help='输入CSV文件路径')
    parser.add_argument('-o', '--output', default='3d_view.html',
                       help='输出HTML文件路径 (默认: 3d_view.html)')
    parser.add_argument('--lon-col', type=int, default=0,
                       help='经度列索引 (默认: 0)')
    parser.add_argument('--lat-col', type=int, default=1,
                       help='纬度列索引 (默认: 1)')
    parser.add_argument('--alt-col', type=int, default=3,
                       help='融合高程列索引 (默认: 3)')
    parser.add_argument('--delimiter', default=None,
                       help='CSV分隔符 (默认: 自动检测)')
    parser.add_argument('--title', default='3D轨迹可视化',
                       help='图表标题 (默认: 3D轨迹可视化)')
    parser.add_argument('--sheet-name', default=0,
                       help='Excel工作表名称或索引 (默认: 0)')

    args = parser.parse_args()

    success = main(
        input_file=args.input_file,
        output_file=args.output,
        lon_col=args.lon_col,
        lat_col=args.lat_col,
        alt_col=args.alt_col,
        delimiter=args.delimiter,
        title=args.title,
        sheet_name=args.sheet_name
    )

    if not success:
        sys.exit(1)

# 示例用法
if __name__ == "__main__":
    # 检查是否有命令行参数
    if len(sys.argv) > 1:
        # 使用命令行接口
        cli_main()
    else:
        # 直接调用示例
        print("3D可视化工具使用示例")
        print("=" * 50)

        # 示例1: 处理Excel文件
        print("\n示例1: 处理Excel数据文件")
        excel_file = r"C:\Users\HL\Documents\xwechat_files\wxid_rdrswaol1qsr21_737e\msg\file\2025-12\08_51_15_间隔高度10m.xlsx"

        if Path(excel_file).exists():
            success = main(
                input_file=excel_file,
                output_file='./10m_3d_view.html',
                title='无人机轨迹3D可视化示例'
            )
        else:
            print(f"示例文件不存在: {excel_file}")
            print("请提供正确的Excel文件路径")
            success = False

        if success:
            print("\n示例代码:")
            print("""
# 直接调用示例
from 3Dview import main

# 处理Excel文件
main('data.xlsx', sheet_name=0)  # 指定工作表

# 处理CSV文件
main('data.csv')

# 自定义参数
main(
    input_file='trajectory.xlsx',
    output_file='flight_path.html',
    lon_col=0,      # 经度列索引
    lat_col=1,      # 纬度列索引
    alt_col=3,      # 融合高程列索引
    sheet_name=0,   # Excel工作表
    title='飞行轨迹可视化'
)
            """)