import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d import Axes3D

# 天体参数设置（简化模型）
AU = 1.5e8  # 天文单位（公里）
earth_orbital_radius = 1.0 * AU
mars_orbital_radius = 1.5 * AU
orbital_speed = 2e4  # 简化轨道速度（km/s）

# 时间参数
earth_period = 365  # 天
mars_period = 687  # 天
transfer_time = 258  # 霍曼转移时间（天）
time_step = 2  # 动画时间步长（天）
total_duration = 800  # 总任务时间（天）

# 初始化图形
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')
ax.set_facecolor('black')


# 天体初始位置
def celestial_pos(t, period, radius):
    angle = 2 * np.pi * t / period
    return radius * np.array([np.cos(angle), np.sin(angle), 0])


# 霍曼转移轨道计算
def transfer_orbit(t, t_start, radius_from, radius_to, transfer_time):
    angle = np.pi * (t - t_start) / transfer_time
    r = radius_from + (radius_to - radius_from) * (t - t_start) / transfer_time
    return r * np.array([np.cos(angle), np.sin(angle), 0])


# 初始化绘图元素
earth, = ax.plot([], [], [], 'o', color='blue', markersize=8)
mars, = ax.plot([], [], [], 'o', color='red', markersize=6)
ship, = ax.plot([], [], [], 'o', color='white', markersize=4)
trajectory, = ax.plot([], [], [], '-', color='gray', alpha=0.5)
sun = ax.plot([0], [0], [0], 'o', color='yellow', markersize=12)[0]

# 坐标轴设置
max_orbit = mars_orbital_radius * 1.2
ax.set_xlim(-max_orbit, max_orbit)
ax.set_ylim(-max_orbit, max_orbit)
ax.set_zlim(-max_orbit / 10, max_orbit / 10)
ax.axis('off')


# 动画更新函数
def update(frame):
    t = frame * time_step

    # 更新天体位置
    earth_pos = celestial_pos(t, earth_period, earth_orbital_radius)
    earth.set_data(earth_pos[0], earth_pos[1])
    earth.set_3d_properties(0)

    mars_pos = celestial_pos(t, mars_period, mars_orbital_radius)
    mars.set_data(mars_pos[0], mars_pos[1])
    mars.set_3d_properties(0)

    # 飞船状态机
    if t < transfer_time:  # 前往火星阶段
        ship_pos = transfer_orbit(t, 0, earth_orbital_radius, mars_orbital_radius, transfer_time)
    elif t < transfer_time + 30:  # 火星停留
        ship_pos = mars_pos
    else:  # 返回地球阶段
        ship_pos = transfer_orbit(t - transfer_time - 30, 0, mars_orbital_radius, earth_orbital_radius, transfer_time)

    ship.set_data(ship_pos[0], ship_pos[1])
    ship.set_3d_properties(0)

    # 更新轨迹
    x, y, z = trajectory.get_data_3d()
    x = np.append(x, ship_pos[0])
    y = np.append(y, ship_pos[1])
    z = np.append(z, 0)
    trajectory.set_data(x, y)
    trajectory.set_3d_properties(z)

    return earth, mars, ship, trajectory, sun


# 创建动画
ani = FuncAnimation(fig, update, frames=int(total_duration / time_step),
                    interval=50, blit=True)

# 添加图例和标注
ax.text(0, 0, 0, "SUN", color='yellow', ha='center')
ax.text(earth_orbital_radius, 0, 0, "Earth", color='blue')
ax.text(mars_orbital_radius, 0, 0, "Mars", color='red')

plt.show()