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 + mid[0], tmp.y + mid[1]), (self.pos.x + mid[0], self.pos.y + mid[1]))

    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, offset, width=1):
        pygame.draw.circle(screen, (0, 0, 0), (int(self.pos.x + mid[0] - offset[0]), int(self.pos.y + mid[1] - offset[1])), self.radius, 0)
        pygame.draw.circle(screen, self.color, (int(self.pos.x + mid[0]- offset[0]), int(self.pos.y + mid[1] - offset[1])), self.radius, width)

running = True
speed = 1

def getEvents():
    for event in pygame.event.get():
        global moving
        if event.type == KEYDOWN:
            global speed, pos, mid, size
            if event.key == pygame.K_UP:
                moving.y = -32
            if event.key == pygame.K_DOWN:
                moving.y = 32
            if event.key == pygame.K_LEFT:
                moving.x = -32
            if event.key == pygame.K_RIGHT:
                moving.x = 32
            if event.key == pygame.K_HOME:
                pos = Rect(mid[0] - (size[0] / 2), mid[1] - (size[1] / 2), size[0], size[1])
            if event.key == ord(']'):
                if speed < 64:
                    speed *= 2
                else:
                    print('There''s a cap at x64!')
            if event.key == ord('['):
                if speed > 2:
                    speed /= 2
                else:
                    print('No pausing or going backwards!')
            if event.key == ord('q'):
                running = False
                sys.exit()
        if event.type == KEYUP:
            if event.key == pygame.K_UP:
                moving.y = 0
            if event.key == pygame.K_DOWN:
                moving.y = 0
            if event.key == pygame.K_LEFT:
                moving.x = 0
            if event.key == pygame.K_RIGHT:
                moving.x = 0
        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(random.randint(10**3, 10**6), pos, Point(getVelocity(pos.y, 10**11), getVelocity(pos.x, 10**11)), 3, randomCol())

pygame.init()
universe = (6000, 4000)
mid = (3000, 2000)
size = (800, 600)
pos = Rect(mid[0] - (size[0] / 2), mid[1] - (size[1] / 2), size[0], size[1])
moving = Point(0, 0)
screen = pygame.display.set_mode(size)
back = pygame.Surface(universe)
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(0, 0), 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(x, y)))

while running:
    getEvents()
    pos[0] += moving.x
    pos[1] += moving.y
    for i in range(0, speed):
        for x in range(0, len(planets)):
            planets[x].resetForce()
        for x in range(0, len(planets)):
            for y in range(x + 1, len(planets)):
                f = getGravity(planets[x], planets[y])
                planets[x].applyForce(f)
                planets[y].applyForce(invertForce(f))
        for x in range(0, len(planets)):
            planets[x].update()

    screen.blit(back, (0, 0), pos)
    for planet in planets:
        planet.draw(screen, pos, 2)

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