# FlockingDispalyDivergence.py
# Run as FlockingDisplayDivergence.py filename

# This program retrieves the time series data of bird locations of an original flock and perturbed copy of the flock
# stored in the files filename & filenamecopy and creates an animation of the bird movements using pygame, 
# to allow visualization of how small differences in inital conditions can lead to large differences in 'final' positions. 

# determine filename
import sys
filename = sys.argv[1]
filenamecopy = filename + 'copy'

# 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 blue 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
      

# initialize the windwow and background to a white screen
screen = pygame.display.set_mode((GS,GS))
background = pygame.Surface([GS, GS])
background.fill(white)
screen.blit(background, [0, 0])
pygame.display.update()


# Get the position data for the Original 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]

# Get the position data for the 'Perturbed' Flock
DataCopy = GetPositions(filenamecopy)
flock_positions_x_copy = DataCopy[0]
flock_positions_y_copy = DataCopy[1]
bird_positions_x_copy = DataCopy[2]
bird_positions_y_copy = DataCopy[3]


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

# initialize a group of bird sprites for the perturbed flock
perturbed_flock = pygame.sprite.RenderUpdates()
for n in range(N):
    perturbed_flock.add( BirdSprite(n, red, [bird_positions_x_copy[n][0]*scalefactor, bird_positions_y_copy[n][0]*scalefactor]) )


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

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

    original_flock.update(x_positions,y_positions)       
    rectlist = original_flock.draw(screen)
    pygame.display.update(rectlist)
    original_flock.clear(screen,background)

    perturbed_flock.update(x_positions_copy,y_positions_copy)       
    rectlist = perturbed_flock.draw(screen)
    pygame.display.update(rectlist)
    perturbed_flock.clear(screen,background)
    #if mod(t,10) == 0: time.sleep(0.003)

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