import pygame
import sys
import math

# --- 配置常量 ---
SCREEN_WIDTH, SCREEN_HEIGHT = 1000, 700
FPS = 60
G_CONSTANT = 2.5  # 自定义引力常数（用于游戏内平衡显示效果）

# 颜色定义
COLOR_BG = (5, 5, 15)         # 深空黑
COLOR_SUN = (255, 200, 50)    # 恒星黄
COLOR_EARTH = (50, 100, 255)  # 行星蓝
COLOR_MOON = (200, 200, 200)  # 卫星灰
COLOR_SHIP = (0, 255, 150)    # 飞船绿
COLOR_TRAIL = (255, 255, 255, 30) # 轨迹白(半透明需特殊处理，这里简化为浅色)

pygame.init()
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Pygame Space Simulator - Gravity & Orbit")
clock = pygame.time.Clock()
font = pygame.font.Font(None, 36)

# --- 基础天体类 (Celestial Body) ---
class CelestialBody(pygame.sprite.Sprite):
    def __init__(self, mass, x, y, vx, vy, color, radius):
        super().__init__()
        self.mass = mass
        self.x, self.y = float(x), float(y)
        self.vx, self.vy = float(vx), float(vy)
        self.color = color
        self.radius = radius
        
        self.image = pygame.Surface((radius*2, radius*2), pygame.SRCALPHA)
        pygame.draw.circle(self.image, color, (radius, radius), radius)
        self.rect = self.image.get_rect(center=(self.x, self.y))

    def apply_gravity(self, other_bodies):
        """根据万有引力定律计算受到的合力并更新速度"""
        total_fx, total_fy = 0.0, 0.0
        
        for body in other_bodies:
            if body is not self:
                dx = body.x - self.x
                dy = body.y - self.y
                dist_sq = max(dx**2 + dy**2, 100)  # 防止距离过近导致除零爆炸
                dist = math.sqrt(dist_sq)
                
                # F = G * (m1 * m2) / r^2
                force_mag = G_CONSTANT * (self.mass * body.mass) / dist_sq
                
                # 分解力到 X 和 Y 轴
                total_fx += force_mag * (dx / dist)
                total_fy += force_mag * (dy / dist)
        
        # a = F / m
        ax = total_fx / self.mass
        ay = total_fy / self.mass
        
        # 欧拉积分更新速度和位置
        self.vx += ax
        self.vy += ay
        self.x += self.vx
        self.y += self.vy
        
        self.rect.center = (int(self.x), int(self.y))

    def draw(self, surface):
        surface.blit(self.image, self.rect)

# --- 玩家飞船类 ---
class Spaceship(pygame.sprite.Sprite):
    def __init__(self, x, y):
        super().__init__()
        self.mass = 10.0
        self.x, self.y = float(x), float(y)
        self.vx, self.vy = 2.0, 1.5  # 初始切向速度，保证能绕圈
        self.thrust_power = 0.15
        self.angle = 0
        
        # 简单的三角形飞船
        self.image = pygame.Surface((40, 40), pygame.SRCALPHA)
        points = [(20, 0), (0, 40), (40, 40)]
        pygame.draw.polygon(self.image, COLOR_SHIP, points)
        self.rect = self.image.get_rect(center=(self.x, self.y))
        
        self.trail = []  # 记录轨迹坐标

    def handle_input(self):
        keys = pygame.key.get_pressed()
        thrust_x, thrust_y = 0.0, 0.0
        
        # WASD 控制推进器方向
        if keys[pygame.K_w]: thrust_y -= self.thrust_power
        if keys[pygame.K_s]: thrust_y += self.thrust_power
        if keys[pygame.K_a]: thrust_x -= self.thrust_power
        if keys[pygame.K_d]: thrust_x += self.thrust_power
        
        # 将推力转换为加速度 (a = F/m)
        self.vx += thrust_x / self.mass
        self.vy += thrust_y / self.mass
        
        # 计算朝向角度 (用于UI或未来加火焰特效)
        if thrust_x != 0 or thrust_y != 0:
            self.angle = math.degrees(math.atan2(-thrust_y, thrust_x))

    def apply_gravity(self, celestial_bodies):
        """飞船也会受到天体引力影响"""
        total_fx, total_fy = 0.0, 0.0
        for body in celestial_bodies:
            dx = body.x - self.x
            dy = body.y - self.y
            dist_sq = max(dx**2 + dy**2, 100)
            dist = math.sqrt(dist_sq)
            
            force_mag = G_CONSTANT * (self.mass * body.mass) / dist_sq
            total_fx += force_mag * (dx / dist)
            total_fy += force_mag * (dy / dist)
            
        self.vx += (total_fx / self.mass)
        self.vy += (total_fy / self.mass)

    def update(self):
        self.x += self.vx
        self.y += self.vy
        self.rect.center = (int(self.x), int(self.y))
        
        # 记录轨迹
        self.trail.append((int(self.x), int(self.y)))
        if len(self.trail) > 300:  # 限制轨迹长度
            self.trail.pop(0)

    def draw(self, surface):
        # 绘制轨迹
        if len(self.trail) > 1:
            pygame.draw.lines(surface, (100, 200, 255), False, self.trail, 2)
        # 绘制飞船
        surface.blit(self.image, self.rect)

# --- 主循环 ---
def main():
    # 创建天体系统 (质量, X, Y, X速度, Y速度, 颜色, 半径)
    sun = CelestialBody(mass=8000, x=500, y=350, vx=0, vy=0, color=COLOR_SUN, radius=30)
    earth = CelestialBody(mass=400, x=750, y=350, vx=0, vy=-4.5, color=COLOR_EARTH, radius=15)
    moon = CelestialBody(mass=10, x=790, y=350, vx=0, vy=-6.5, color=COLOR_MOON, radius=5)
    
    celestial_bodies = [sun, earth, moon]
    all_sprites = pygame.sprite.Group(celestial_bodies)
    
    ship = Spaceship(300, 350)
    
    running = True
    while running:
        clock.tick(FPS)
        
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            if event.type == pygame.KEYDOWN and event.key == pygame.K_r:
                main()  # 简单粗暴的重置方式
                return

        # 物理引擎更新
        for body in celestial_bodies:
            body.apply_gravity(celestial_bodies)
            
        ship.handle_input()
        ship.apply_gravity(celestial_bodies)
        ship.update()
        
        # 绘制
        screen.fill(COLOR_BG)
        all_sprites.draw(screen)
        for body in celestial_bodies:
            body.draw(screen)
        ship.draw(screen)
        
        # UI提示
        hint_text = font.render("WASD Thrust | R Reset", True, (150, 150, 150))
        vel_text = font.render(f"Velocity: {math.hypot(ship.vx, ship.vy):.2f}", True, COLOR_SHIP)
        screen.blit(hint_text, (10, SCREEN_HEIGHT - 30))
        screen.blit(vel_text, (10, 10))

        pygame.display.flip()

    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()