# FlockingDisplayObstacles.py
# Run as FlockingDisplayObstacles.py filename Obstacle_file

# This program retrieves the time series data of bird locations, stored in the file filename,
# and creates an animation of the bird movements using pygame. It also displays the obstacles
# stored in the the file Obstacle_file which the birds had to negotiate in the simulation
# (this may be an empty textfile but the program still needs it as input to run). 

# determine filename
import sys
filename = sys.argv[1]
Obstacle_file = sys.argv[2]

# import modules
from FlockingGetPositions import *
from numpy import *
import time 
import pygame
pygame.init()

# define some colors
white = 255,255,255
black = 0,0,0
blue = 0,0,255
red = 255,0,0

# Define Some Parameters 
# gridsize = the size of the grid the original data was simulated on
# scalefactor = the factor by which this data is scaled up for display purposes
gridsize = 100
scalefactor = 8.0
GS = gridsize*scalefactor


# define the class 'BirdSprite'. They are displayed as boxes, and have their own update rule to move positions.
class BirdSprite(pygame.sprite.Sprite):
    def __init__(self, number, color, position):
        pygame.sprite.Sprite.__init__(self)
        self.number = number
        self.image = pygame.Surface([scalefactor, scalefactor])
        self.image.fill(color)
        self.rect = self.image.get_rect()
        self.rect.topleft = position

    def update(self, x_positions, y_positions):
        position = [ x_positions[self.number], y_positions[self.number] ]
        self.rect.topleft = position
      

# Get the positions of obstacles
Obstacles = []
f = open(Obstacle_file)
for line in f:
    l = line[:-1]
    ll = l.split()
    Obstacles.append([ scalefactor*float(ll[0]), scalefactor*float(ll[1]) ])


# initialize the windwow and create a background (white screen with obstacles)
screen = pygame.display.set_mode((GS,GS))
background = pygame.Surface([GS, GS])
background.fill(white)

for o in Obstacles:
    TempSurface = pygame.Surface([scalefactor, scalefactor])
    TempSurface.fill(black)
    background.blit(TempSurface, o)

screen.blit(background, [0, 0])
pygame.display.update()


# Get the position data for the Flock
Data = GetPositions(filename)
flock_positions_x = Data[0]
flock_positions_y = Data[1]
bird_positions_x = Data[2]
bird_positions_y = Data[3]
N = Data[4]
num_timesteps = Data[5]


# initialize a group of bird sprites for the flock
flock = pygame.sprite.RenderUpdates()
for n in range(N):
    flock.add( BirdSprite(n, blue, [bird_positions_x[n][0]*scalefactor, bird_positions_y[n][0]*scalefactor]) )


# run the loop to display moving BirdSprites
for t in xrange(1,num_timesteps):
    
    x_positions = []
    y_positions = []

    for n in xrange(N):
        x_positions.append(bird_positions_x[n][t]*scalefactor)
        y_positions.append(bird_positions_y[n][t]*scalefactor)

    flock.update(x_positions,y_positions)       
    rectlist = flock.draw(screen)
    pygame.display.update(rectlist)
    flock.clear(screen,background)
    #if mod(t,10) == 0: time.sleep(0.003)   (Use this command to slow down the display if it is too fast)

    if t == 1: time.sleep(3.0)