"""
Holding the left mouse button will repel birds from wherever you clicked, the right one attracts them. Use the little red x to close the window, I've found that Python crashes when I just close the window.

If you have any questions or suggestions or anything else: http://philippdow.blogspot.com/
"""


import sys
from math import *
import random
import pygame
from pygame.locals import *

pygame.init()
zeroview=[]
running = True                                      #is the script running? gets set to False when the small x gets clicked, and the program then stops.

#colors
green = (0,200,0)
white=(255,255,255)
black=(0,0,0)
red=(250,0,0)


#general variables
screensize = [1000,700]                             #screen size in pixels
loopscreen = False                                  #Lets the screen loop, although radius calculations then don't expand to the opposite side of the screen
screen = pygame.display.set_mode((screensize[0],screensize[1]))
pygame.draw.rect(screen, white, (0,0, screensize[0], screensize[1]),0)


numbirds=80                         #number of birds. around 60-70 is good, but lower numbers allow the program to run much faster

birds=[["bird"+str(i),[0,0],[0,0],[0,0],[0,0]]
           for i in range(numbirds)]                #the list of birds, in the form (name, position, velocity, last velocity, second last velocity)*number of birds

radius=20                                           #maximum distance of two birds to be considered "neighbours". standard: 50
viewangle=abs(cos(1.65806279))                      #angle a bird can see, to each side of the current direction vector, around 95 is standard. the cos() is so that the isinview(x,y) function doesn't need to call acos() for every bird every frame


averagebirdspeed=9                                  #average speed of the birds in pixels/frame. standard: 8
lowerbound,upperbound=3,3                           #this lets the user control how far the speed of the birds can deviate - speed is going to be random in the range of averagebirdspeed-lowerbound to averagebirdspeed+upperbound. 3 and 3 are standard.



#view options for visualization:
connectbirds=False                                  #interconnect all birds (very slow if there are many birds)
connectaverageposition=False                        #connects all birds to the swarms average position
showaveragevector=False                             #shows the swarms average velocity vector (length is proportional to magnitude, but the factor is not necessarily 1)
connectaveragetocenter=False                        #draws a line from the center to the birds average position
highlightbird=False                                 #highlights a single bird to visualize its trajectory
traceimage=False                                    #if true, the white background doesn't get redrawn, and the bird trajectories will be traced out. also traces any lines that are being drawn.


#weightings for velocity components:
randomweight=20                                     #25
neighbourweight=0                                   #10
neighbourposweight=0                                #0.5
mouseweight=200                                     #200, must be high for mouse attraction/repulsion to work properly
viewweight=16                                       #15
viewposweight=1                                     #1

#weighting for last velocity components, higher values here will create smoother flight curves
o1weight=50                                         #50
o2weight=15                                         #40

#gravity, to keep the birds from leaving the visible frame for too long
gravstrength = 0.0004                               #1/2500 is good. This basically decides how far out of the screen area the birds will fly.
gravitypoint=(
        [int(screensize[0]/2),int(screensize[1]/2)])#middle of screensize is best






#small functions, mostly for making it easier to work with the vectors
def plusminus():                                    #random 1 or -1
    if (random.random()>0.5):
        return 1
    else: return -1

def vmag(v):                                        #returns a vectors magnitude
    return abs(sqrt( v[0]**2 + v[1]**2 ))

def addvectorcomponents(a,b):                       #adds the components of a and b
    return [a[0]+b[0], a[1]+b[1]]

def subvectorcomponents(a,b):                       #subtracts the components of a and b
    return [a[0]-b[0], a[1]-b[1]]

def averagevector(vectors):                         # accepts a list of vectors as input, returns the average vector
    xsum=0
    ysum=0
    for i in range(len(vectors)):
        xsum+=vectors[i][0]
        ysum+=vectors[i][1]
    return [xsum/len(vectors), ysum/len(vectors)]

def scalarmultivector(x,vect):                      #multiply vect by the scalar x
    return [x*vect[0], x*vect[1]]

def randomspeedvector():                            #generates a random speed vector from the variables averagebirdspeed, upper and lower bound
    return [random.randrange(1,averagebirdspeed*2+1)*plusminus(), random.randrange(1,averagebirdspeed*2+1)*plusminus()]

def scalevector(vector, magnitude):                 #scale a vector to a certain magnitude without altering its direction
    ratio=magnitude/vmag(vector)
    return [vector[0]*ratio, vector[1]*ratio]




#more specific functions for calculations regarding the bird's positions and velocities
def isinview(x, y):                                 #can bird x see bird y?

    if x==y:                                        #bird doesn't see itself
        return False

    b1b2v = [birds[x][1][0]-birds[y][1][0],         #vector from bird 1 to bird 2
             birds[x][1][1]-birds[y][1][1]]
    if ((( (b1b2v[0]*birds[x][2][0] + b1b2v[1]*birds[x][2][1])/(vmag(b1b2v)*vmag(birds[x][2]))))) <= viewangle:
        return True                                 #iff the angle between the birds velocity vector and the vector from bird 1 to bird 2 is within the birds viewing angle, bird 1 can see bird 2

    else:
        return False

def updatevelocity(b):
    global running
    global zeroview
    bird=birds[b]
    neighbours=[[],[]]                              #list for neighbour velocity and position calculations
    view=[[],[]]                                    #list for calculation of average velocity of birds within view
    birds[b][4]=birds[b][3]                         #update previous velocities
    birds[b][3]=birds[b][2]


    if neighbourweight and neighbourposweight:
        for x in range(numbirds):                   #this block calculates the neighbour birds' average velocity and the neighbour birds' average position
            if (vmag(subvectorcomponents(birds[b][1], birds[x][1]))<=radius):
                neighbours[0].append(birds[x][2])
                neighbours[1].append(birds[x][1])
        neighbouraveragevelocity=averagevector(neighbours[0])
        neighbouraveragepos=subvectorcomponents(averagevector(neighbours[1]),birds[b][1])
    else:
        neighbouraveragevelocity=[0,0]
        neighbouraveragepos=[0,0]


    if viewweight:
        for x in range(numbirds):                   #this block calculates the average velocity of the visible birds
            if isinview(b,x):
                view[0].append(birds[x][2])
                view[1].append(birds[x][1])
        if len(view[0]):
            viewaveragevelocity=averagevector(view[0])
            viewaverageposition=subvectorcomponents(averagevector(view[1]),birds[b][1])
        else:
            viewaveragevelocity=[0,0]
            viewaverageposition=[0,0]
    else:
        viewaveragevelocity=[0,0]
        viewaverageposition=[0,0]
    if b==0:
        zeroview=view[1][:]

    randomcomponent=randomspeedvector()             #add a random component, the birds "free will"

    gravweight=(vmag(subvectorcomponents(gravitypoint, birds[b][1])))
    gravitycomponent=subvectorcomponents(gravitypoint, birds[b][1])



    mouse=pygame.mouse.get_pos()                    #this if-block checks whether the mouse is focused on the window and pressed, and then either attracts or repels birds from that point depending on which button is pressed

    if pygame.mouse.get_focused() and pygame.mouse.get_pressed()[0]:

        if ((mouse[0] >= (screensize[0]-15))        #checks if the "close" button has been pressed
            and (mouse[1] <= (15))):
            running = False

        bmV=subvectorcomponents(birds[b][1], mouse)
        mousecomponent=scalevector(bmV, 1)
        mousestrength=1/vmag(mousecomponent)**2

    elif pygame.mouse.get_focused() and pygame.mouse.get_pressed()[2]:

        bmV=subvectorcomponents(mouse,birds[b][1])
        mousecomponent=scalevector(bmV, 1)
        mousestrength=1/vmag(mousecomponent)**2

    else:
        bmV=subvectorcomponents(mouse,birds[b][1])
        mousecomponent=scalevector(bmV, 1)
        mousestrength=0

    #update the actual velocity, taking into account the weightings

    birds[b][2] = [((neighbourweight*neighbouraveragevelocity[0]+
                        viewweight*viewaveragevelocity[0]+
                        viewposweight*viewaverageposition[0]+
                        randomweight*randomcomponent[0]+
                        neighbourposweight*neighbouraveragepos[0]+
                        gravstrength*gravweight*gravitycomponent[0]+
                        mouseweight*mousestrength*mousecomponent[0]+
                        o1weight*birds[b][3][0]+
                        o2weight*birds[b][4][0])/
                    (viewposweight+randomweight+neighbourweight+neighbourposweight+viewweight+o1weight+o2weight+gravstrength*gravweight)),


                        ((neighbourweight*neighbouraveragevelocity[1]+
                        viewweight*viewaveragevelocity[1]+
                        viewposweight*viewaverageposition[1]+
                        randomweight*randomcomponent[1]+
                        neighbourposweight*neighbouraveragepos[1]+
                        gravstrength*gravweight*gravitycomponent[1]+
                        mouseweight*mousestrength*mousecomponent[1]+
                        o1weight*birds[b][3][1]+
                        o2weight*birds[b][4][1])/
                    (viewposweight+randomweight+neighbourweight+neighbourposweight+viewweight+o1weight+o2weight+gravstrength*gravweight))]

    #Finally, we adjust the speed to a random speed in the range of averagebirdspeed-lowerbound and averagebirdspeed+upperbound. Not crucial, but makes the simulation look more natural (depending on how variables are set)

    birds[b][2]=scalevector(birds[b][2], random.randrange(averagebirdspeed-lowerbound,averagebirdspeed+upperbound))


def updateposition(b):
    if loopscreen:                                  #buggy for view and neighbour calculations
        birds[b][1]=[(birds[b][1][0]+birds[b][2][0])%screensize[0], (birds[b][1][1]+birds[b][2][1])%screensize[1]]

    else:                                           #for screen that doesn't loop
        birds[b][1]=[(birds[b][1][0]+birds[b][2][0]), (birds[b][1][1]+birds[b][2][1])]

def averagebirdposition():
    return averagevector([birds[b][1] for b in range(numbirds)])

def averagebirdvelocity():
    vel=[ birds[b][2] for b in range(numbirds)]
    return scalevector(averagevector(vel), vmag(averagevector(vel))*8)





# "system" functions
def setup():                                        #initializes birds with random positions and velocities
    for b in range(numbirds):
        birds[b][2] = birds[b][3] = birds[b][4] = randomspeedvector()
        birds[b][1] = [random.randrange(0, screensize[0]), random.randrange(0, screensize[1])]


def refreshscreen():                                #draw the birds, also includes conditionals for the various visualization options
    if not traceimage:
        pygame.draw.rect(screen, white, (0,0, screensize[0], screensize[1]),0)

    count=0
    for bird in birds:
        pygame.draw.rect(screen, black, (bird[1][0], bird[1][1], 2, 2), 0)
        #if bird[0]=="bird0":
         #   pygame.draw.rect(screen, green, (bird[1][0], bird[1][1], 4, 4), 0)
        if highlightbird and count==0:              #this bit can highlight bird 0 to demonstrate a specific bird's motion
           pygame.draw.rect(screen, red, (bird[1][0], bird[1][1], 3, 3), 1)

        if connectbirds:
            for b in birds:
                pygame.draw.line(screen, green, bird[1], b[1], 1)
        if connectaverageposition:
            pygame.draw.line(screen, black, averagebirdposition(), bird[1], 1)
    #for zv in zeroview:
       # pygame.draw.rect(screen, red, (zv[0], zv[1], 3, 3), 0)

    if connectaveragetocenter:
        pygame.draw.line(screen, green, [screensize[0]/2,screensize[1]/2], averagebirdposition(), 1)
    if showaveragevector:
        pygame.draw.line(screen, red, averagebirdposition(), [averagebirdvelocity()[0]+averagebirdposition()[0],averagebirdvelocity()[1]+averagebirdposition()[1]], 2)

    pygame.draw.rect(screen, red, (screensize[0]-16 , 0, 16,16),0)
    pygame.draw.line(screen, black, [screensize[0]-15,15], [screensize[0],0], 1)
    pygame.draw.line(screen, black, [screensize[0],14], [screensize[0]-15,0], 1)


    pygame.display.update()




setup()
#main program loop
while True:
    pygame.event.pump()

    for b in range(numbirds):
        updatevelocity(b)
        updateposition(b)

    if running == False:
        pygame.quit()
        sys.exit()

    refreshscreen()