from Numeric import *
import random
from pylab import *
import pygame
import getopt,sys
from numpy import *
from visual.graph import *
from visual import *
from Scientific.Statistics import average, variance, standardDeviation
from Scientific.IO.TextFile import TextFile

try:
       import Numeric as N
       import pygame
       import pygame.surfarray as surfarray
except ImportError:
       raise ImportError, "Numeric and PyGame required."



#User needs to tell the program what file to read in and the dimensions
file=TextFile('lattice1.txt', 'r')
N=128
M=128





row=0
col=0
curstate=zeros((int(128),int(128)))
nextstate=zeros((N,M))
curstate1=transpose(curstate)


#print curstate
for line in file:
    col=0
    #print len(line)
    for element in line:
        
        if element != "\n":
            
          #  print element,
            
            curstate[row][col]=int(element)
    
            #print curstate[row][col]
            col=col+1
        
    row=row+1





count=0.
for i in range (M):
    for j in range (N):
        if curstate[i][j]==1 or curstate[i][j]==2:
            count=count+1.
#print "count="
#print count
        

CellSize = 4




empty=0,0,0
southbound=255,0,0
eastbound=0,0,255
size=(N*CellSize,M*CellSize)
screen = pygame.display.set_mode(size)
win=500
L = 1.
v=array([0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.])
avg=array([0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.])
x=zeros(100)
for i in range(100):
    x[i]=i
    

curstate=transpose(curstate)




def iterate(t):
    q=0.    
     
    if t%2==0:
    #print "move 2"
    #MOVE THE TWOS DOWN
        
        
        for i in range (N):
            for j in range (M):
                
                    if curstate[i][j]==2 and curstate[(i+1)%M][j]==0:
            			nextstate[(i+1)%M][j]=2
            			nextstate[i][j]=0
            			q=q+1.
                    elif curstate[i][j]==2:
                        nextstate[i][j]=2
                    if curstate[i][j]==1:
                        nextstate[i][j]=1
            	
            	
    if t%2==1:  
    #MOVE THE ONES RIGHT
      #  print "move 1"

        
        for i in range (N):
            for j in range (M):
                
                    
                    
                
                    if curstate[i][j]==1 and curstate[i][(j-1)%N]==0:
                        nextstate[i][(j-1)%N]=1
                    	nextstate[i][j]=0
                    	q=q+1.
                    elif curstate[i][j]==1:
                        nextstate[i][j]=1
                    if curstate[i][j]==2:
                        nextstate[i][j]=2
                
	
    curstate[:][:]=nextstate[:][:]
    return q
    
    
def animate(state,nSites,mSites,surface,CellSize):
       ConfigRect = pygame.Rect(0,0,nSites*CellSize,mSites*CellSize)
       surface.fill(empty,ConfigRect)
       for i in range(nSites):
               for j in range(mSites):
                        if state[i][j] == 1:
                            surface.fill(eastbound,[i*CellSize,j*CellSize,CellSize,CellSize])
                        if state[i][j] == 2:
                            surface.fill(southbound,[i*CellSize,j*CellSize,CellSize,CellSize])
                        
#    print t
#    print curstate

pygame.display.set_caption('2D Traffic Cellular Automaton Simulator')
       # Clear display
screen.fill(empty)
pygame.display.flip()
sptmdiag = surfarray.pixels3d(screen)





t=0
pygame.display.flip()
while t>-1:
    print t
    for event in pygame.event.get():
               if event.type == pygame.QUIT: sys.exit()
    if t%2==0:
        q1=iterate(t)
        
    if t%2==1:
        q2=iterate(t)
        if t!=1:
        	velocity=(q1+q2)/count
        

	        for i in range(99):
	            v[99-i]=v[98-i]
	        v[0]=velocity
        
        #print standardDeviation(v)
        #print variance(v)
        
	        for i in range(99):
	            avg[99-i]=avg[98-i]
	        avg[0]=average(v)

    animate(curstate,N,M,screen,CellSize)
    pygame.display.flip()
    print v[99]
    
    t=t+1