import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d.art3d import Line3DCollection
import imageio

def spherical_to_cartesian(lat, lon, r):
    lat_rad = np.radians(lat)
    lon_rad = np.radians(lon)
    x = r * np.cos(lat_rad) * np.cos(lon_rad)
    y = r * np.cos(lat_rad) * np.sin(lon_rad)
    z = r * np.sin(lat_rad)
    return x, y, z

def interpolate_on_sphere(point1, point2, interval, r):
    lat1, lon1 = point1
    lat2, lon2 = point2
    
    # 计算插值点数
    num_points = int(np.ceil(np.hypot(lat2 - lat1, lon2 - lon1) / interval))

    if num_points < 2:
        num_points = 2

    t = np.linspace(0, 1, num_points)
    lat = (1 - t) * lat1 + t * lat2
    lon = (1 - t) * lon1 + t * lon2
    
    # 将插值点转换为笛卡尔坐标，并加上偏移量
    cartesian_points = [spherical_to_cartesian(lat_i, lon_i, r) for lat_i, lon_i in zip(lat, lon)]
    
    return cartesian_points

def plot_trajectory_on_sphere(trajectory, start, end, radius=1, offset=0.05, interval=5, obstacles=None):
    """
    在球面上绘制一条轨迹线，并标出起点和终点，轨迹线浮在球面上方。

    参数:
    trajectory (list of tuples): 轨迹点列表，每个点为 (经度, 纬度) 形式。
    start (tuple): 起始位置 (经度, 纬度)。
    end (tuple): 目标位置 (经度, 纬度)。
    radius (float): 球面半径，默认为 1。
    offset (float): 轨迹线与球面之间的偏移量，默认为 0.05。
    interval (int): 每段轨迹之间的插值间隔，默认为 5 度。
    obstacles (list of tuples): 障碍区域列表，每个区域为 (经度, 纬度, 半径) 形式。
    """

    # 将轨迹点转换为笛卡尔坐标，并加上偏移量
    interpolated_trajectory = []
    for i in range(len(trajectory) - 1):
        interpolated_trajectory.extend(interpolate_on_sphere(trajectory[i], trajectory[i + 1], interval, radius))
    
    # 添加最后一个点
    interpolated_trajectory.append(spherical_to_cartesian(*trajectory[-1], radius))

    # 生成球面数据
    u = np.linspace(0, 2 * np.pi, 100)
    v = np.linspace(0, np.pi, 100)
    u, v = np.meshgrid(u, v)
    x = radius * np.outer(np.cos(u), np.sin(v))
    y = radius * np.outer(np.sin(u), np.sin(v))
    z = radius * np.outer(np.ones(np.size(u)), np.cos(v))

    # 初始化颜色数组为蓝色
    colors = np.full((*x.shape, 4), [0, 0, 1, 1])

    # 标记障碍区域的颜色
    if obstacles is not None:
        for obstacle in obstacles:
            lat, lon, r_obstacle = obstacle
            center_x, center_y, center_z = spherical_to_cartesian(lat, lon, radius)

            # 生成障碍区域的网格
            mask = ((x - center_x)**2 + (y - center_y)**2 + (z - center_z)**2) <= r_obstacle**2
            colors[mask] = [1, 1, 0, 0.1]  # 黄色，透明度为 0.1

    # 创建3D图形
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # 生成更细密的三角面片
    u, v = np.mgrid[0:2 * np.pi:100j, 0:np.pi:100j]
    x = radius * np.cos(u) * np.sin(v)
    y = radius * np.sin(u) * np.sin(v)
    z = radius * np.cos(v)

    # 生成三角形索引
    triangles = []
    for i in range(99):
        for j in range(99):
            p1 = i * 100 + j
            p2 = i * 100 + j + 1
            p3 = (i + 1) * 100 + j
            p4 = (i + 1) * 100 + j + 1
            triangles.extend([(p1, p2, p3), (p2, p3, p4)])

    # 绘制球面
    ax.plot_trisurf(x.flatten(), y.flatten(), z.flatten(), triangles=triangles, facecolors=colors.flatten(), shade=False)

    # 标记起点和终点
    start_x, start_y, start_z = spherical_to_cartesian(*start, radius + offset)
    end_x, end_y, end_z = spherical_to_cartesian(*end, radius + offset)
    ax.scatter(start_x, start_y, start_z, color='r', s=100, label='Start')
    ax.scatter(end_x, end_y, end_z, color='g', s=100, label='End')

    # 设置坐标轴标签
    ax.set_xlabel('X axis')
    ax.set_ylabel('Y axis')
    ax.set_zlabel('Z axis')

    # 添加图例
    ax.legend()

    # 生成动图 GIF
    images = []
    for i in range(len(interpolated_trajectory)):
        # 清除之前的轨迹线
        for collection in ax.collections:
            del collection
        
        # 绘制部分轨迹线
        trajectory_x, trajectory_y, trajectory_z = zip(*interpolated_trajectory[:i+1])
        points = np.array(list(zip(trajectory_x, trajectory_y, trajectory_z))).reshape(-1, 1, 3)
        segments = np.concatenate([points[:-1], points[1:]], axis=1)
        lc = Line3DCollection(segments, cmap='viridis', linewidth=2)
        lc.set_array(np.linspace(0, 1, len(trajectory_x)))
        ax.add_collection3d(lc)

        # 获取图像数据
        canvas = fig.canvas
        canvas.draw()
        image = np.frombuffer(canvas.buffer_rgba(), dtype='uint8')
        image = image.reshape(fig.canvas.get_width_height()[::-1] + (4,))
        images.append(image[:, :, :3])  # 只取 RGB 部分

    imageio.mimsave('trajectory.gif', images, duration=0.1)

    plt.show()  # 显示图形
    plt.close(fig)  # 关闭图形

# 示例用法
trajectory = [(0, 0), (30, 30), (60, 60), (90, 90)]
start = (0, 0)
end = (90, 90)
radius = 1
offset = 0.05
interval = 5  # 每段轨迹之间的插值间隔
obstacles = [(45, 45, 0.1), (75, 75, 0.1)]  # 障碍区域列表

plot_trajectory_on_sphere(trajectory, start, end, radius, offset, interval, obstacles)