mixster

#!/usr/bin/python
import sys, pygame
from pygame.locals import *
import math, random

class Point(object):
    def __init__(self, x=0, y=0):
        self.x = x
        self.y = y

class Body(object):
    def __init__(self, mass=1, pos=(0, 0), vel=Point(0, 0), radius=0, color=(0, 0, 0)):
        self.pos = pos
        self.radius = radius
        self.acc = Point(0, 0)
        self.vel = vel
        self.mass = mass
        self.color = color

    def move(self, dx, dy):
        tmp = self.pos
        self.pos.x += dx
        self.pos.y += dy
        pygame.draw.line(back, self.color, (tmp.x, tmp.y), (self.pos.x, self.pos.y))

    def resetForce(self):
        self.acc = Point(0, 0)

    def applyForce(self, force):
        self.acc.x += force.x / self.mass
        self.acc.y += force.y / self.mass

    def update(self):
        self.vel.x += self.acc.x
        self.vel.y += self.acc.y
        self.move(self.vel.x, self.vel.y)

    def draw(self, screen, width=1):
        pygame.draw.circle(screen, self.color, (int(self.pos.x), int(self.pos.y)), self.radius, width)

running = True
speed = 1

def getEvents():
    for event in pygame.event.get():
        if event.type == KEYDOWN:
            global speed
            if event.key == K_UP:
                if speed < 64:
                    speed *= 2
                else:
                    print('There''s a cap at x64!')
            if event.key == K_DOWN:
                if speed > 2:
                    speed /= 2
                else:
                    print('No pausing or going backwards!')
            if event.key == ord('q'):
                running = False
                sys.exit()
        if event.type == pygame.QUIT:
            running = False
            sys.exit()

def getGravity(bodyA, bodyB):
    dx = bodyA.pos.x - bodyB.pos.x
    dy = bodyA.pos.y - bodyB.pos.y
    d = dx**2 + dy**2

    if d == 0:
        return Point(0, 0)
    else:
        f = ((6.673*(10**-11)) * bodyA.mass * bodyB.mass) / d
        a = math.atan2(dy, dx)
        fx = f * -math.cos(a)
        fy = f * -math.sin(a)

        return Point(fx, fy)

def randomCol():
    x = random.randint(0, 255)
    y = random.randint(0, 255)
    z = random.randint(0, 255)
    if x + y + z >= 100:
        return (x, y, z)
    else:
        return randomCol()

def invertForce(force):
    return Point(-force.x, -force.y)

def getVelocity(r, mass):
    if r != 0:
        return abs((6.673*(10**-11)*mass)/r)**0.5*(random.randrange(-1, 2, 2))
    else:
        return 0

def getDefaultBody(pos):
    return Body(1, pos, Point(getVelocity(200 - pos.y, 10**11), getVelocity(300 - pos.x, 10**11)), 5, randomCol())

pygame.init()
screen = pygame.display.set_mode((600, 400))
back = pygame.Surface((600, 400))
back.fill((0, 0, 0))
pygame.display.set_caption('Gravity simulator')
pygame.mouse.set_visible(0)

clock = pygame.time.Clock()
random.seed()

planets = []
planets.append(Body(10**11, Point(300, 200), Point(0, 0), 2, (0, 0, 255)))

for i in range(1, 20):
    x, y = 0, 0
    if random.randint(0, 1) == 0:
        x = random.randint(25, 200)
    else:
        y = random.randint(25, 200)

    if random.randint(0, 1) == 0:
        x = -x
        y = -y

    planets.append(getDefaultBody(Point(300 + x, 200 + y)))

while running:
    getEvents()
    for i in range(0, speed):
        for planet in planets:
            planet.resetForce()
            for otherplanet in planets:
                planet.applyForce(getGravity(planet, otherplanet))
            planet.update()

    #screen.fill((0, 0, 0))
    screen.blit(back, (0, 0))
    #pygame.draw.rect(screen, (0, 0, 0), pygame.Rect(0, 0, 600, 400))
    for planet in planets:
        planet.draw(screen, 2)

    clock.tick(30)
    pygame.display.flip()
print('Terminating')