"""
质点类，表示三体问题中的一个天体
"""

import numpy as np
from typing import Tuple, List


class Particle:
    """表示一个质点的类"""
    
    def __init__(self, mass: float, position: np.ndarray, velocity: np.ndarray, 
                 name: str = "", color: str = None):
        """
        初始化质点
        
        参数:
            mass: 质量（太阳质量）
            position: 位置向量 (x, y, z) [AU]
            velocity: 速度向量 (vx, vy, vz) [AU/年]
            name: 质点名称
            color: 可视化颜色
        """
        self.mass = mass
        self.position = np.array(position, dtype=np.float64)
        self.velocity = np.array(velocity, dtype=np.float64)
        self.name = name
        self.color = color
        
        # 存储轨迹历史
        self.position_history = []
        self.velocity_history = []
        
    def update(self, new_position: np.ndarray, new_velocity: np.ndarray):
        """更新质点的状态并记录历史"""
        self.position_history.append(self.position.copy())
        self.velocity_history.append(self.velocity.copy())
        self.position = new_position
        self.velocity = new_velocity
    
    def get_trajectory(self) -> np.ndarray:
        """获取轨迹历史"""
        if not self.position_history:
            return np.array([self.position])
        return np.array(self.position_history + [self.position])
    
    def get_energy(self) -> float:
        """计算质点的动能"""
        v_squared = np.sum(self.velocity**2)
        return 0.5 * self.mass * v_squared
    
    def __repr__(self) -> str:
        return (f"Particle(name='{self.name}', mass={self.mass:.3f}, "
                f"position={self.position}, velocity={self.velocity})")
    
    def copy(self) -> 'Particle':
        """创建质点的副本"""
        return Particle(
            mass=self.mass,
            position=self.position.copy(),
            velocity=self.velocity.copy(),
            name=self.name,
            color=self.color
        )