# -*- coding:utf-8 -*-
# title           :二维天体运行模拟
# description     :二维天体运行模拟，适合初学者学习。
# author          :Python超人
# date            :2025-03-23
# link            :https://gitcode.net/pythoncr/
# python_version  :3.8
# ==============================================================================
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# 设置 matplotlib 支持中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体字体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 定义物理常量
# 万有引力常量，单位：m^3 kg^-1 s^-2
万有引力常量 = 6.67430e-11


# 定义星球类
class 星球:
    def __init__(self, 质量, 位置, 速度, 颜色, 大小=None, 名称=None, 名称偏移=[0, 0], 轨迹最大长度=2000):
        """
        初始化星球对象。

        :param 质量: 星球的质量，单位：kg
        :param 位置: 星球的初始位置，是一个二维列表或数组，单位：m
        :param 速度: 星球的初始速度，是一个二维列表或数组，单位：m/s
        :param 颜色: 星球的颜色，用于绘图显示
        :param 大小: 星球在图中的显示大小，若无指定则根据质量计算
        :param 名称: 星球的名称，用于在图中显示
        :param 名称偏移: 星球名称显示位置相对于星球位置的偏移量，是一个二维列表或数组，单位：m
        :param 轨迹最大长度: 星球轨迹的最大记录长度，无单位
        """
        if 大小 is None:
            # 大小 = np.sqrt(质量) / 1000
            大小 = np.power(质量, 0.2) / 10
        self.质量 = 质量
        self.位置 = np.array(位置, dtype=float)
        self.速度 = np.array(速度, dtype=float)
        self.颜色 = 颜色
        self.大小 = 大小
        self.名称 = 名称
        self.名称偏移 = 名称偏移
        self.轨迹 = [self.位置.copy()]
        self.轨迹最大长度 = 轨迹最大长度

    def 更新(self, 星球列表, 时间步长):
        """
        更新星球的位置和速度。

        :param 星球列表: 包含所有星球对象的列表
        :param 时间步长: 每次更新的时间间隔，单位：s
        """
        总力 = np.zeros(2)
        for 星球 in 星球列表:
            if 星球 is not self:
                距离向量 = 星球.位置 - self.位置
                距离 = np.linalg.norm(距离向量)
                # 计算万有引力，力的单位：N（kg m s^-2）
                力 = 万有引力常量 * self.质量 * 星球.质量 * 距离向量 / 距离 ** 3
                总力 += 力
        # 计算加速度，单位：m s^-2
        加速度 = 总力 / self.质量
        # 更新速度，单位：m/s
        self.速度 += 加速度 * 时间步长
        # 更新位置，单位：m
        self.位置 += self.速度 * 时间步长
        self.轨迹.append(self.位置.copy())
        if len(self.轨迹) > self.轨迹最大长度:
            self.轨迹.pop(0)


# 创建星球列表
星球列表 = [
    星球(质量=1e12, 位置=[0, 0], 速度=[0, 1], 颜色='red', 名称="太阳", 名称偏移=[2, 0]),
    星球(质量=1e9, 位置=[10, 0], 速度=[0, 2], 颜色='blue', 名称="蓝星"),
    星球(质量=1e9, 位置=[-10, 0], 速度=[0, -2], 颜色='green', 名称="绿星")
]

# 初始化图形
图形, 坐标轴 = plt.subplots()
坐标轴.set_xlim(-50, 50)
坐标轴.set_ylim(-20, 80)
# 调整边距
plt.subplots_adjust(left=0.05, bottom=0.04, right=0.97, top=0.97, wspace=0.4, hspace=0.4)
点列表 = [坐标轴.plot([], [], 'o', color=星球.颜色, markersize=星球.大小)[0] for 星球 in 星球列表]
轨迹列表 = [坐标轴.plot([], [], '-', color=星球.颜色)[0] for 星球 in 星球列表]
名称列表 = [坐标轴.text(0, 0, 星球.名称, color=星球.颜色) for 星球 in 星球列表]


# 初始化函数
def 初始化():
    """
    初始化动画的状态。
    """
    for 点 in 点列表:
        点.set_data([], [])
    for 轨迹 in 轨迹列表:
        轨迹.set_data([], [])
    for 名称 in 名称列表:
        名称.set_position((0, 0))
    return 点列表 + 轨迹列表 + 名称列表


# 更新函数
def 更新(帧):
    """
    更新每一帧动画中星球的位置和轨迹。

    :param 帧: 当前动画的帧数，无单位
    """
    # 时间步长，单位：s
    时间步长 = 0.01
    for 星球 in 星球列表:
        星球.更新(星球列表, 时间步长)
    for i, 星球 in enumerate(星球列表):
        点列表[i].set_data(星球.位置[0], 星球.位置[1])
        轨迹 = np.array(星球.轨迹)
        轨迹列表[i].set_data(轨迹[:, 0], 轨迹[:, 1])
        偏移量_x, 偏移量_y = 星球.名称偏移
        名称列表[i].set_position((星球.位置[0] + 偏移量_x + 1, 星球.位置[1] + 偏移量_y))
    return 点列表 + 轨迹列表 + 名称列表


# 创建动画
动画 = FuncAnimation(图形, 更新, frames=range(1000), init_func=初始化, blit=True, interval=2)

# 显示动画
plt.show()