main.py

import pygame
import random
import sys
import math

pygame.init()
WIDTH, HEIGHT = 800, 600
SCREEN = pygame.display.set_mode((WIDTH, HEIGHT))
CLOCK = pygame.time.Clock()

G = 100  # gravitational constant (tweak for visuals)
DT = 0.01  # time step

class Body:
    def __init__(self, x, y, vx, vy, m, color):
        self.pos = pygame.math.Vector2(x, y)
        self.vel = pygame.math.Vector2(vx, vy)
        self.m = m
        self.color = color

    def draw(self):
        radius = max(4, int(math.log(self.m + 1) * 2))
        pygame.draw.circle(SCREEN, self.color, (int(self.pos.x), int(self.pos.y)), radius)

def compute_acc(bodies, i_idx):
    acc = pygame.math.Vector2(0, 0)
    bi = bodies[i_idx]
    for j, bj in enumerate(bodies):
        if i_idx == j: continue
        r_vec = bj.pos - bi.pos
        dist_sq = r_vec.length_squared() + 100  # Softens close-range forces
        acc += G * bj.m * r_vec / (dist_sq * math.sqrt(dist_sq))
    return acc

def step_euler(bodies):
    accs = [compute_acc(bodies, i) for i in range(len(bodies))]
    for b, a in zip(bodies, accs):
        b.pos += b.vel * DT
        b.vel += a * DT

def step_rk4(bodies):
    n = len(bodies)
    # Save initial states
    pos0 = [b.pos.copy() for b in bodies]
    vel0 = [b.vel.copy() for b in bodies]

    # k1
    a1 = [compute_acc(bodies, i) for i in range(n)]

    # k2
    for i, b in enumerate(bodies):
        b.pos = pos0[i] + vel0[i] * (DT / 2)
        b.vel = vel0[i] + a1[i] * (DT / 2)
    a2 = [compute_acc(bodies, i) for i in range(n)]

    # k3
    for i, b in enumerate(bodies):
        b.pos = pos0[i] + b.vel * (DT / 2)
        b.vel = vel0[i] + a2[i] * (DT / 2)
    a3 = [compute_acc(bodies, i) for i in range(n)]

    # k4
    for i, b in enumerate(bodies):
        b.pos = pos0[i] + b.vel * DT
        b.vel = vel0[i] + a3[i] * DT
    a4 = [compute_acc(bodies, i) for i in range(n)]

    # Combine
    for i, b in enumerate(bodies):
        b.pos = pos0[i] + vel0[i] * DT + (DT**2 / 6) * (a1[i] + 2*a2[i] + 2*a3[i] + a4[i])
        b.vel = vel0[i] + (DT / 6) * (a1[i] + 2*a2[i] + 2*a3[i] + a4[i])

def random_color():
    return (random.randint(100, 255), random.randint(100, 255), random.randint(100, 255))

def create_bodies(n=20):
    bodies = []
    for _ in range(n):
        x = random.uniform(100, WIDTH - 100)
        y = random.uniform(100, HEIGHT - 100)
        vx = random.uniform(-50, 50)
        vy = random.uniform(-50, 50)
        m = random.uniform(100, 200)
        bodies.append(Body(x, y, vx, vy, m, random_color()))
    return bodies

def create_solar_system():
    sun = Body(WIDTH // 2, HEIGHT // 2, 0, 0, 5000, (255, 255, 0))  # Fixed central mass

    orbiters = []
    num_orbiters = 8
    for i in range(num_orbiters):
        angle = i * (2 * math.pi / num_orbiters)
        dist = random.uniform(80, 250)
        x = WIDTH // 2 + dist * math.cos(angle)
        y = HEIGHT // 2 + dist * math.sin(angle)

        # Tangential velocity for circular orbit
        dir_to_center = pygame.math.Vector2(WIDTH // 2 - x, HEIGHT // 2 - y)
        tangential = pygame.math.Vector2(-dir_to_center.y, dir_to_center.x).normalize()
        speed = math.sqrt(G * sun.m / dist)
        vx, vy = tangential * speed

        m = random.uniform(50, 300)
        orbiters.append(Body(x, y, vx, vy, m, random_color()))

    return [sun] + orbiters

bodies = create_solar_system()
initial_bodies = [(b.pos.copy(), b.vel.copy()) for b in bodies]

integrator = 'Euler'  # default
running = False

while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_e:
                integrator = 'Euler'
            elif event.key == pygame.K_r:
                integrator = 'RK4'
            elif event.key == pygame.K_SPACE:
                bodies = create_solar_system()
            elif event.key == pygame.K_p:
                running = not running

    SCREEN.fill((0, 0, 0))

    if running:
        if integrator == 'Euler':
            step_euler(bodies)
        else:
            step_rk4(bodies)

    for b in bodies:
        b.draw()

    # Display mode
    font = pygame.font.SysFont(None, 24)
    text = font.render(f'Integrator: {integrator}', True, (255, 255, 255))
    SCREEN.blit(text, (10, 10))

    pygame.display.flip()
    CLOCK.tick(60)