2D ball bouncing collisions code

import pygame, sys, random, math, time, os
from pygame.locals import *

bif = "burst.jpg"
mif = "red-dot.png"
ball1image = "ball.png"
ball2image = "golf-ball.png"
ball3image = "circle.png"
boing = "flop.wav"


WIDTH = 1025
HEIGHT = 822
SPEED = 1
count = 0
x_reflect_count = 0
y_reflect_count = 0

def load_sound(name):
    class NoneSound:
        def play(self): pass
    if not pygame.mixer:
        return NoneSound()
    try:
        sound = pygame.mixer.Sound(name)
    except pygame.error, message:
        print 'Cannot load sound:', name
        raise SystemExit, message
    return sound

pygame.init()
pygame.mixer.pre_init(44100, -16, 2, 2048)
pygame.mixer.init()

screen = pygame.display.set_mode((WIDTH,HEIGHT),0,32)

background = pygame.image.load(bif).convert()
WIDTH = background.get_width()
HEIGHT = background.get_height()

cursor = pygame.image.load(mif).convert_alpha()
cursor_radius = cursor.get_width()/2

ball1 = pygame.image.load(ball1image).convert_alpha()
ball1_radius = ball1.get_width()/2

class Ball(object):
    def __init__(self, image, idnum = None):
        """pos and vel are tuples; image is a path to the image file"""
        self.x = 0 # type float
        self.y = 0
        self.vx = 0
        self.vy = 0
        self.image = pygame.image.load(image).convert_alpha() # type pygame image object
        self.radius = self.image.get_width()/2
        self.x_refcount = 0
        self.y_refcount = 0
        self.count = 0
        self.idnum = idnum
        self.mass = math.pi * 4/3 * self.radius**3
    def get_position(self):
        """Return's the ball's position vector (center of ball)"""
        return (self.x, self.y)
    def get_velocity(self):
        """Return's the ball's velocity vector"""
        return (self.vx, self.vy)
    def get_speed(self):
        return math.sqrt(self.vx ** 2 + self.vy ** 2)
    def get_direction(self):
        return math.atan2(self.vy, self.vx)
    def get_radius(self):
        """Returns the ball's radius"""
        return self.radius
    def get_draw_coords(self):
        """Returns the coordinates that the ball will be drawn at"""
        x_draw = self.x - self.radius
        y_draw = self.y - self.radius
        return (x_draw, y_draw)
    def get_edges(self):
        """Return left, right, top and bottom edges"""
        return [self.x - self.radius, self.x + self.radius, self.y - self.radius, self.y + self.radius]
    def move(self):
        """Updates position inx by moving it delta, delta is a float"""
        self.x = (self.x + self.vx * SPEED) % (WIDTH + 1 + SPEED)
        self.y = (self.y + self.vy * SPEED) % (HEIGHT + 1 + SPEED)
    def update_vel(self, vel):
        """Set the velocity to arbitrary vel tuple"""
        self.vx = vel[0]
        self.vy = vel[1]
    def update_pos(self, pos):
        """Set the position to arbitrary pos tuple"""
        self.x = pos[0]
        self.y = pos[1]
    def reverse_x(self):
        """Flip x direction"""
        self.vx = -1 * self.vx
    def reverse_y(self):
        """Flip y direction"""
        self.vy = -1 * self.vy
    def draw(self, disp, coords = None):
        """Draw the ball on the screen, disp is a pygame display object, coords is a tuple"""
        if not coords: coords = self.get_draw_coords()
        disp.blit(self.image, coords)
    def update_x_refcount(self, increment = True):
        if increment:
            self.x_refcount += 1
        else:
            self.x_refcount = 0
    def update_y_refcount(self, increment = True):
        if increment:
            self.y_refcount += 1
        else:
            self.y_refcount = 0
    def get_x_refcount(self):
        return self.x_refcount
    def get_y_refcount(self):
        return self.y_refcount
    def get_idnum(self):
        return self.idnum
    def get_mass(self):
        return self.mass
    def get_count(self):
        return self.count
    def update_count(self, increment = True):
        if increment:
            self.count += 1
        else:
            self.count = 0

def set_inc(theta = None):
    if not theta:
        theta = 2 * random.random() * math.pi
    x_inc = math.cos(theta)
    y_inc = math.sin(theta)
    return x_inc, y_inc


def manage_bounce(ball1, ball2):
    ball1_coords = ball1.get_position()
    ball1_vel = ball1.get_velocity()
    ball2_coords = ball2.get_position()
    ball2_vel = ball2.get_velocity()

    dx = ball1_coords[0] - ball2_coords[0]
    dy = ball1_coords[1] - ball2_coords[1]
    collision_angle = math.atan2(dy, dx)
    magnitude_1 = ball1.get_speed()
    magnitude_2 = ball2.get_speed()
    direction_1 = ball1.get_direction()
    direction_2 = ball2.get_direction()

    new_xspeed_1 = magnitude_1 * math.cos(direction_1 - collision_angle);
    new_yspeed_1 = magnitude_1 * math.sin(direction_1 - collision_angle);
    new_xspeed_2 = magnitude_2 * math.cos(direction_2 - collision_angle);
    new_yspeed_2 = magnitude_2 * math.sin(direction_2 - collision_angle);

    final_xspeed_1 = ((ball1.get_mass() - ball2.get_mass()) * new_xspeed_1 + 2 * ball2.get_mass() * new_xspeed_2) / (ball1.get_mass() + ball2.get_mass())
    final_xspeed_2 = (2 * ball.get_mass() * new_xspeed_1 + (ball2.get_mass() - ball1.get_mass()) * new_xspeed_2) / (ball1.get_mass() + ball2.get_mass())
    final_yspeed_1 = new_yspeed_1;
    final_yspeed_2 = new_yspeed_2;

    ball1_xspeed = math.cos(collision_angle)*final_xspeed_1 + math.cos(collision_angle + math.pi / 2) * final_yspeed_1
    ball1_yspeed = math.sin(collision_angle)*final_xspeed_1 + math.sin(collision_angle + math.pi / 2) * final_yspeed_1
    ball1.update_vel((ball1_xspeed, ball1_yspeed))
    ball2_xspeed = math.cos(collision_angle)*final_xspeed_2 + math.cos(collision_angle + math.pi / 2) * final_yspeed_2
    ball2_yspeed = math.sin(collision_angle)*final_xspeed_2 + math.sin(collision_angle + math.pi / 2) * final_yspeed_2
    ball2.update_vel((ball2_xspeed, ball2_yspeed))


ball1 = Ball(ball1image, 1)
ball1.update_pos((WIDTH/4, HEIGHT/2))
ball1.update_vel(set_inc())

ball2 = Ball(ball2image, 2)
ball2.update_pos((3*WIDTH/4, HEIGHT/2))
ball2.update_vel(set_inc())

ball3 = Ball(ball3image, 3)
ball3.update_pos((WIDTH/2, 3*HEIGHT/4))
ball3.update_vel(set_inc())

balls = [ball1, ball2, ball3]


pillar1 = Ball(mif, 1)
pillar1.update_pos((WIDTH/4, HEIGHT/4))
pillar2 = Ball(mif, 2)
pillar2.update_pos((WIDTH/4, 3*HEIGHT/4))
pillar3 = Ball(mif, 3)
pillar3.update_pos((3*WIDTH/4, HEIGHT/4))
pillar4 = Ball(mif, 4)
pillar4.update_pos((3*WIDTH/4, 3*HEIGHT/4))
pillar5 = Ball(mif, 5)
pillar5.update_pos((WIDTH/2, HEIGHT/2))

pillars = [pillar1, pillar2, pillar3, pillar4, pillar5]


cursor_oop = Ball(mif)

ball1_x, ball1_y = WIDTH/2,HEIGHT/2
x_inc, y_inc = set_inc()

font = pygame.font.Font(None, 25)
bounce_sound = load_sound(boing)

pygame.mixer.init()

frame_end = time.clock()

while True:

    for event in pygame.event.get():
        if event.type == QUIT:
            pygame.quit()
            sys.exit()
        if event.type == KEYDOWN:
            if event.key == K_UP:
                if SPEED < 0.1: SPEED += 0.01
                else: SPEED += 0.1
            elif event.key == K_DOWN:
                if SPEED <= 0.1: SPEED -= 0.01
                else: SPEED -= 0.1
            else:
                ball1.update_vel(set_inc())
            print "speed = ", SPEED

    screen.blit(background, (0,0))

    cursor_oop.update_pos(pygame.mouse.get_pos())

    #moving and wall collisions#
    for ball in balls:
        ball.move()
        ball_edges = ball.get_edges()
        if ball_edges[0] <= 0 or ball_edges[1] >= WIDTH:
            if ball.get_x_refcount() > 10:
                ball.reverse_x()
                ball.update_x_refcount(False)
            x_reflect_count = 0
        if ball_edges[2] <= 0 or ball_edges[3] >= HEIGHT:
            if ball.get_y_refcount() > 10:
                ball.reverse_y()
                ball.update_y_refcount(False)

    cursor_oop.draw(screen)
    for ball in balls:
        ball.draw(screen)
    for pillar in pillars:
        pillar.draw(screen)

    #ball-ball collisions#
    compared_pairs = {(ball, ball) for ball in balls} # we will keep track of which pairs we have compared here
    for ballA in balls:
        for ballB in balls:
            if (ballA, ballB) not in compared_pairs:
                ballA_coords = ballA.get_position()
                ballB_coords = ballB.get_position()
                dist_x = ballB_coords[0] - ballA_coords[0]
                dist_y = ballB_coords[1] - ballA_coords[1]
                dist_abs = math.sqrt(dist_x**2 + dist_y**2)
                if (dist_abs < ballA.get_radius() + ballB.get_radius()) and ballA.get_count() > 10 and ballB.get_count() > 10:
                    print "collision between ball %d and ball %d!" % (ballA.get_idnum(), ballB.get_idnum())
                    ballA.update_count(False)
                    ballB.update_count(False)
                    manage_bounce(ballA, ballB)
                compared_pairs.add((ballA.get_idnum(),ballB.get_idnum())) # add pair to set so it won't be compared again
                compared_pairs.add((ballB.get_idnum(),ballA.get_idnum()))

        for pillar in pillars:
            dist_x = pillar.get_position()[0] - ball.get_position()[0]
            dist_y = pillar.get_position()[1] - ball.get_position()[1]
            dist_abs = math.sqrt(dist_x**2 + dist_y**2)
            if (dist_abs < pillar.get_radius() + ball.get_radius()):
                if True:#ball.get_count() > 10:
                    print "ball %d hit pillar %d!" % (ball.get_idnum(),pillar.get_idnum())
                    ball.update_count(False)
                    ball.update_vel((-dist_x/dist_abs, -dist_y/dist_abs))
                else: print "count was too low (%d)" % ball.get_count()

    #mouse-ball collisions#
    cursor_coords = cursor_oop.get_position()
    for ball in balls:
        ball_coords = ball.get_position()
        dist_x = cursor_coords[0] - ball_coords[0]
        dist_y = cursor_coords[1] - ball_coords[1]
        dist_abs = math.sqrt(dist_x**2 + dist_y**2)

        if (dist_abs < cursor_oop.get_radius() + ball.get_radius()) and count > 180:
            print "hit ball %d! it took you %d" % (ball.get_idnum(),count)
            count = 0
            ball.update_vel((-dist_x/dist_abs, -dist_y/dist_abs))

    #measuring framerate#
    previous_frame = frame_end
    frame_end = time.clock()

    elapsed_time = frame_end-previous_frame
    fps = 1/elapsed_time

    #displaying text on screen#
    dist_text = font.render("Dist x: %0.02f Dist y: %0.02f Dist abs: %0.02f" % (dist_x, dist_y, dist_abs), 1, (10,10,10))
    screen.blit(dist_text, (0,0))
    posn_text = font.render("Cursor x: %0.02f Cursor y: %0.02f" % cursor_oop.get_position(), 1, (10,10,10))
    screen.blit(posn_text, (0,15))
    time_text = font.render("Time elapsed: %0.07f FPS: %0.07f" % (elapsed_time, fps), 1, (10,10,10))
    screen.blit(time_text, (0,30))

    #bookkeeping#
    count += 1
    for ball in balls:
        ball.update_x_refcount()
        ball.update_y_refcount()
        ball.update_count()
    pygame.display.update()