import pygame
from pygame.locals import *


class vector:
	def __init__(self, x, y):
		self.x = x
		self.y = y

	def length(self):
		return (self.x ** 2 + self.y ** 2) ** (0.5)

	def norm(self):
		d = self.length()
		self.x /= d
		self.y /= d

	def norm_to(self, l):
		self.norm()
		self.x *= l
		self.y *= l

	def __add__(self, other):
		return vector(self.x + other.x, self.y + other.y)


K = 1


class planet:
	def __init__(self, x, y, r, vx, vy, m=None):
		if m is None:
			m = r ** 3
		self.x = x
		self.y = y
		self.r = r
		self.m = m
		self.speed = vector(vx, vy)

	def move(self):
		self.x += self.speed.x
		self.y += self.speed.y

	def draw(self, surface):
		pygame.draw.circle(surface, (255, 0, 0), (round(self.x), round(self.y)), self.r)

	def gravity_pull(self, other):
		res = vector(other.x - self.x, other.y - self.y)
		l = K * self.m * other.m / dist_squared(self, other)
		res.norm_to(l)
		return res


def dist_squared(first, second):
	return (first.x - second.x) ** 2 + (first.y - second.y) ** 2


width = 1800
height = 900

pygame.init()
display = pygame.display.set_mode((width, height))
clock = pygame.time.Clock()
FPS = 30

obj_list = []

obj_list.append(planet(width / 2, height / 2 - 400, 20, -18, 0))
obj_list.append(planet(width / 2, height / 2 + 200, 10, 20, 0))
obj_list.append(planet(width / 2, height / 2, 50, 0, 0))

# obj_list.append(planet(width / 2 - 500, height / 2 + 300, 40, 35, 0))
# obj_list.append(planet(width / 2 - 500, height / 2 + 150, 5, 55, 0))
# obj_list.append(planet(width / 2, 0, 100, -3, 0))

while True:
	display.fill((0, 0, 0))
	for obj in obj_list:
		obj.draw(display)
	pygame.display.update()

	for obj in obj_list:
		obj.move()

	for obj1 in obj_list:
		Fravn = vector(0, 0)
		for obj2 in obj_list:
			if obj1 is obj2:
				continue
			Fravn = Fravn + obj1.gravity_pull(obj2)
		acc = vector(Fravn.x / obj1.m, Fravn.y / obj1.m)
		obj1.speed = obj1.speed + acc

	clock.tick(FPS)
	for event in pygame.event.get():
		if event.type == QUIT:
			pygame.quit()
			exit()