import pygame
import sys
import random
import objects
import math

class Particle():
    def __init__(self, x, y, size, color):
        self.x = x
        self.y = y
        self.size = size
        self.colour = color
        self.thickness = 1
        self.speed = 2
        self.angle = math.pi
        self.drag = 0.999
        self.elasticity = 0.75
        self.selected = False

    def bounce(self, width, height):
        left_border = 0
        rightborder = width
        top = 0
        bottom = height

        if self.x+self.size >= rightborder:
            #self.speed = 0
            self.x = 2 * (width - self.size) - self.x
            self.angle = - self.angle
            self.speed *= self.elasticity

        if self.x-self.size <= left_border:
            #self.speed = 0
            self.x = 2 * self.size - self.x
            self.angle = - self.angle
            self.speed *= self.elasticity

        if self.y+self.size >= bottom:
            #self.speed = 0
            self.y = 2 * (height - self.size) - self.y
            self.angle = math.pi - self.angle
            self.speed *= self.elasticity

        if self.y-self.size <= top:
            #self.speed = 0
            self.y = 2 * self.size - self.y
            self.angle = math.pi - self.angle
            self.speed *= self.elasticity

        

    def move(self):
        self.x += math.sin(self.angle) * self.speed
        self.y -= math.cos(self.angle) * self.speed
        
        # Gravity
        self.angle, self.speed = addVectors(self.angle, self.speed, math.pi, 0.02)
        
        # Drag - if im understanding it ast the ball moves its losing 0.001 of its speed
        drag = 0.999
        self.speed *= drag

    def display(self, surface):
        pygame.draw.circle(surface, self.colour, (self.x, self.y), self.size)

class EntityObject():
    def __init__(self, parent, x, y):
        self.radius = 32
        self.parent = parent
        self.color = (255, 255, 255)
        self.surface = pygame.Surface((64, 64), pygame.SRCALPHA, 32)
        self.rect = pygame.draw.circle(self.surface, self.color, (self.radius, self.radius), self.radius)

        self.position = objects.Vector(x, y)

    def update(self):
        pass

    def show(self):
        self.parent.blit(self.surface, (self.position.x, self.position.y))


def addVectors(angle1, length1, angle2, length2):
    x = math.sin(angle1) * length1 + math.sin(angle2) * length2
    y = math.cos(angle1) * length1 + math.cos(angle2) * length2
    length = math.hypot(x, y)
    angle = 0.5 * math.pi - math.atan2(y, x)
    return (angle, length)

class Game():
    def __init__(self):
        super(Game, self).__init__()
        self.height = 760
        self.width = 1280
        self.window = pygame.display.set_mode((self.width, self.height))
        self.window_rect = self.window.get_rect()
        self.clock = pygame.time.Clock()
        self.window.set_colorkey((35, 35, 35))
        self.debug_grid = True
        self.background = (35, 35, 35)

        self.particles = []

        n_particles = 30

        for _ in range(n_particles):
            size = random.randint(20, 40)
            red = random.randint(100, 255)
            blue = random.randint(100, 255)
            green = random.randint(100, 255)
            x = random.randint(size, self.width-size)
            y = random.randint(size, self.height-size)
            particle = Particle(x, y, size, (red, blue, green))
            particle.speed = random.random() + 2
            particle.angle = random.uniform(0, math.pi*2)
            self.particles.append(particle)

    def main(self):
        self.event_handler()
        self.update()
        self.blit()

    def event_handler(self):
        """
        This method handles mouse and keyboard events for the scene and passes
        them on as is needed.
        """
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit(0)
            if event.type == pygame.MOUSEBUTTONDOWN:
                mouseX, mouseY = pygame.mouse.get_pos()
                selected_particle = self.findParticle(mouseX, mouseY)
                if selected_particle:
                    selected_particle.selected = True
            if event.type == pygame.MOUSEBUTTONUP:
                self.unselectParticles()
                

    def update(self):
        pass

    def unselectParticles(self):
        for particle in self.particles:
            particle.selected = False

    def findParticle(self, mouseX, mouseY):
        for particle in self.particles:
            if math.hypot(particle.x-mouseX, particle.y-mouseY) <= particle.size:
                return particle
        return None

    def collide(self, p1, p2):
        dx = p1.x - p2.x
        dy = p1.y - p2.y
        
        distance = math.hypot(dx, dy)
        if distance < p1.size + p2.size:
            print('Bang!')
            tangent = math.atan2(dy, dx)
            p1.angle = 2 * tangent - p1.angle
            p2.angle = 2 * tangent - p2.angle
            (p1.speed, p2.speed) = (p2.speed, p1.speed)
            p1.speed *= p1.elasticity
            p2.speed *= p2.elasticity
            angle = 0.5 * math.pi + tangent
            p1.x += math.sin(angle)
            p1.y -= math.cos(angle)
            p2.x -= math.sin(angle)
            p2.y += math.cos(angle)

    def blit(self):
        self.window.fill(self.background)
        if self.debug_grid:
            for y in range(1, int(self.window_rect.height/32) + 1):
                pygame.draw.line(self.window, (150, 150, 150), (0, y * 32), (self.window_rect.width, y * 32))

            for x in range(1, int(self.window_rect.width/32) + 1):
                pygame.draw.line(self.window, (150, 150, 150), (x * 32, 0), (x * 32, self.window_rect.height))

        for i, particle in enumerate(self.particles):
            if not particle.selected:
                particle.move()
                particle.bounce(self.width, self.height)
                for particle2 in self.particles[i+1:]:
                    self.collide(particle, particle2)
            else:
                mouseX, mouseY = pygame.mouse.get_pos()
                #particle.x = mouseX
                #particle.y = mouseY
                dx = mouseX - particle.x
                dy = mouseY - particle.y
                particle.angle = math.atan2(dy, dx) + 0.5 * math.pi
                particle.speed = math.hypot(dx, dy) * 0.1
            particle.display(self.window)


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


pygame.display.set_caption('Physics galore!')

pygame.init()
pygame.font.init()
scene = Game()

while scene is not None:
    scene.main()