# Import the pygame library and initialise the game engine
import pygame
import numpy as np
pygame.init()

# Initialize Pygame
size = (800, 800)
CENTER = np.array(size) / 2
WHITE = (255, 255, 255)
DAMPING = 0.99
screen = pygame.display.set_mode(size)
pygame.display.set_caption("Patrick")
running = True
clock = pygame.time.Clock()


class Ball(object):
    def __init__(self, x, y, radius=10, mass=1):
        self.pos = np.array([x, y])
        self.last_pos = np.array([x, y])
        self.velocity = np.array([0, 0])
        self.radius = radius
        self.mass = mass

    def render(self, screen):
        x, y = self.pos + CENTER
        pygame.draw.circle(screen, WHITE, [int(x), int(y)], self.radius, 2)

    def update(self):
        self.pos[1] += .05

        if self.pos[1] > 400:
            self.pos[1] = 400
            self.last_pos[1] = 400
            self.velocity[1] = -.5 * self.velocity[1]

        self.pos += self.velocity

        self.velocity = self.pos - self.last_pos
        self.velocity *= DAMPING
        self.last_pos = np.array(self.pos)


class Spring(object):
    def __init__(self, a, b, length=None, k=0.3):
        self.a = a
        self.b = b
        self.k = k
        self.length = self.compute_length() if length is None else length
        print(self.length)

    def render(self, screen):
        x1, y1 = self.a.pos + CENTER
        x2, y2 = self.b.pos + CENTER
        pygame.draw.line(screen, WHITE, [x1, y1], [x2, y2])

    def compute_length(self):
        x1, y1 = self.a.pos + CENTER
        x2, y2 = self.b.pos + CENTER
        return np.sqrt((x2 - x1)**2 + (y2 - y1)**2)

    def update(self):
        current_length = self.compute_length()
        delta = self.b.pos - self.a.pos
        delta /= current_length
        stretch = (current_length - self.length) / self.length
        self.a.velocity += stretch * delta * self.k
        self.b.velocity -= stretch * delta * self.k


if __name__ == '__main__':
    balls = [Ball(0., 0.), Ball(100., 100.)]
    springs = [Spring(balls[0], balls[1], length=100)]

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        screen.fill((0, 0, 0))

        # Update physics
        for ball in balls:
            ball.update()

        for spring in springs:
            spring.update()

        # Draw
        for ball in balls:
            ball.render(screen)

        for spring in springs:
            spring.render(screen)

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

    pygame.quit()