#! /usr/bin/env python

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

PASTURE = Rect(0, 0, 500, 500)
CULTURALDIST = 5.0
FLOCKSIZE = 50
EXITSSTAGELEFT = False

COMFORTDIST = 5.0
JAYGIRTH = 10.0
WORMGIRTH = 4.0
MAXSPEED = 10.0 # mph, 1 pixel = 1 mile
STRUTSPERHOUR = 2

def render_background():
    surf = pygame.Surface(PASTURE.size)
    surf.fill((255,255,255))
    return surf

def render_rect(size, color):
    surf = pygame.Surface(size)
    r = surf.get_rect()
    surf.fill(color)
    return surf, r

def rand_pos(buffer = 0):
    return random.randrange(buffer, PASTURE.right - buffer), random.randrange(buffer, PASTURE.bottom - buffer)    

class rectangible:
    def __init__(self,x,y,w,h):
        #force float
        self.x = x + 0.; self.y = y + 0.; self.w = w + 0.; self.h = h + 0.
    def colliderect(self, *args):
        if len(args) == 1 and isinstance(args[0], rectangible):
            x = args[0].x; y = args[0].y; w = args[0].w; h = args[0].h
        elif len(args) == 4:
            x = args[0]; y = args[1]; w = args[2]; h = args[3]
        else: raise InputError, "rectangible.colliderect() requires either a rectangible or x,y,w,h"
            
        return ((self.x >= x and self.x < x+w) or (x >= self.x and x < self.x+self.w)) and \
               ((self.y >= y and self.y < y+h) or (y >= self.y and y < self.y+self.h))
    def inflate(self, x, y):
        return rectangible(self.x - x/2, self.y - y/2, self.w + x, self.h + y)
    def move(self, x, y):
        return rectangible(self.x + x, self.y + y, self.w, self.h)
    def move_ip(self, x,y):
        self.x += x; self.y += y
        return self
    def centerx(self):
        return self.x + self.w/2
    def centery(self):
        return self.y + self.h/2
    def Rect(self):
        return pygame.Rect(self.x, self.y, self.w, self.h)
    
class nightcrawler(pygame.sprite.Sprite):
    def __init__(self, size, pos, b):
        pygame.sprite.Sprite.__init__(self, self.containers)
        self.image, self.rect = render_rect(size, (100,70,50))
        self.rect.left, self.rect.top = pos
        self.pos = rectangible(pos[0],pos[1],size[0],size[1])
        self.birthmark = b

class jaybird(pygame.sprite.Sprite):
    def __init__(self, size, pos, m, l, b):
        pygame.sprite.Sprite.__init__(self, self.containers)
        self.image, self.rect = render_rect(size, (0,0,255))
        self.rect.left, self.rect.top = pos
        self.pos = self.lastpos = self.nextpos = rectangible(pos[0],pos[1],size[0],size[1])
        self.maxspeed = m
        self.lunch = l
        self.birthmark = b
        self.sated = False
    def eye(self, fellows):
        xmove = ymove = count = 0
        xdel = self.lunch.centerx() - self.pos.centerx()
        ydel = self.lunch.centery() - self.pos.centery()
        
        for chirper in fellows:
            if chirper != self and not EXITSSTAGELEFT and self.pos.inflate(2*COMFORTDIST, 2*COMFORTDIST).colliderect(chirper.pos):
                count = count + 1
                # move away from the nearby chirper
                xdel2 = self.pos.centerx() - chirper.pos.centerx()
                ydel2 = self.pos.centery() - chirper.pos.centery()
                fiddle2 = (self.maxspeed/2.0)/math.sqrt(xdel2*xdel2 + ydel2*ydel2)
                xmove += fiddle2*xdel2; ymove += fiddle2*ydel2

        # move slightly to the right w.r.t. the direction towards the worm
        fiddle3 = (self.maxspeed/8.0)/math.sqrt(xdel*xdel + ydel*ydel)
        xmove += fiddle3*ydel; ymove += fiddle3*xdel;

       	#move in the direction of the worm, if not in others' midst
        fiddle = (self.maxspeed/4.0)/math.sqrt(xdel*xdel + ydel*ydel)
        xmove += fiddle*xdel; ymove += fiddle*ydel

        return xmove, ymove
    def strut(self):
            #print self.birthmark, ": ", self.pos
            self.lastpos = self.pos
            self.pos = self.nextpos
            self.rect = self.pos.Rect()

def main(winstyle = 0):
    pygame.init()

    # display mode
    winstyle = 0
    bestdepth = pygame.display.mode_ok(PASTURE.size, winstyle, 32)
    screen = pygame.display.set_mode(PASTURE.size, winstyle, bestdepth)

    # background
    bg = render_background()
    screen.blit(bg, (0,0))
    pygame.display.flip()

    # sprite groups
    thebirds = pygame.sprite.Group()
    theworms = pygame.sprite.Group()    
    all = pygame.sprite.RenderUpdates()
    jaybird.containers = thebirds, all
    nightcrawler.containers = theworms, all

    # controls framerate
    clock = pygame.time.Clock()
    
    random.seed()

    # place sprites randomly, but not too close together
    for birthmark in range(FLOCKSIZE):
        while True:
            pos = rand_pos(buffer=JAYGIRTH)
            for wriggler in theworms.sprites():
                if wriggler.pos.inflate(math.sqrt(2)*(WORMGIRTH+2*CULTURALDIST),math.sqrt(2)*(WORMGIRTH+2*CULTURALDIST)).\
                   colliderect(pos[0],pos[1],WORMGIRTH,WORMGIRTH): break
            else: break
        n = nightcrawler((WORMGIRTH, WORMGIRTH), pos, birthmark)
        
        while True:
            pos = rand_pos(buffer=JAYGIRTH/2)
            for chirper in thebirds.sprites():
                if chirper.pos.inflate(CULTURALDIST, CULTURALDIST).colliderect(pos[0], pos[1], JAYGIRTH, JAYGIRTH): break
            else: break
        j = jaybird((JAYGIRTH, JAYGIRTH), pos, MAXSPEED, n.pos, birthmark)


    start = time.time()
    
    # main loop
    flockishly = True
    while flockishly:
	# EVENTS
        pygame.event.pump()
	for event in pygame.event.get():
            if event.type == KEYDOWN and event.key == K_p:
                pygame.event.set_allowed(None)
                pygame.event.set_allowed(KEYDOWN)
                pygame.event.clear()
                pygame.event.set_blocked(None)
                pygame.event.wait()
	    if event.type == QUIT or (event.type == KEYDOWN and event.key == K_q):
		flockishly = False

	# EYE
	for chirper in thebirds.sprites():
            if chirper.sated == False:
                xvel, yvel = chirper.eye(thebirds.sprites()) #XXX: need send only non sated participants if exiting stage left
                chirper.nextpos = chirper.pos.move(xvel/STRUTSPERHOUR, yvel/STRUTSPERHOUR)

	# STRUT
	for chirper in thebirds.sprites():
            if chirper.sated == False:
                chirper.strut()

        # CULTURAL VIOLATIONS & COLLISIONS
	#  uhhh... O(n^2), not O(n!) 
	for chirper in thebirds.sprites():
            if chirper.sated == True: continue
            for chirpier in thebirds.sprites():
		if chirper != chirpier and \
                   chirper.pos.inflate(2*CULTURALDIST, 2*CULTURALDIST).colliderect(chirpier.pos) and \
                   (EXITSSTAGELEFT == False or (EXITSSTAGELEFT == True and chirpier.sated == False)):
                    print chirper.birthmark, " and ", chirpier.birthmark, "are invading each other's personal space!"
                    continue

            	if chirper != chirpier and chirper.pos.colliderect(chirpier.pos) and \
                   (EXITSSTAGELEFT == False or (EXITSSTAGELEFT == True and chirpier.sated == False)):
                    print "collision between ", chirper.birthmark, " and ", chirpier.birthmark
                    print chirper.pos.Rect(), " and ", chirpier.pos.Rect()
                    flockishly = False

        # GOT WORM?
	for chirper in thebirds.sprites():
            if chirper.pos.colliderect(chirper.lunch):
                chirper.sated = True

	# DRAW
	all.clear(screen, bg)
	all.update()
	dirty = all.draw(screen)
        pygame.display.update(dirty)

	# HUNGRYBIRD?
	for chirper in thebirds.sprites():
	    if chirper.sated == False: break
        else: flockishly = False

	#limit framerate
	#clock.tick(100)		

    print "Took ", time.time() - start, "seconds (ignoring pauses, etc.)"

    #pause before exiting
    while True:
        event = pygame.event.wait()
        if event.type == QUIT or event.type == KEYDOWN:
            return

#call the "main" function if running this script
if __name__ == '__main__': main()