from bitarray import bitarray
from random import getrandbits
import pygame
from pygame.locals import *
import numpy as np


def step(mem): 
    
    two03 = ~mem[0] & ~mem[1] & mem[2] & ~mem[3]& ~mem[4] & mem[5]
    two14 = mem[0] & ~mem[1] & ~mem[2] & mem[3]& ~mem[4]& ~mem[5]
    two25 = ~mem[0] & mem[1] & ~mem[2] & ~mem[3]& mem[4]& ~mem[5]
    
    three024 = ~mem[0] &  mem[1] & ~mem[2]  &  mem[3] & ~mem[4]&  mem[5]
    three135 =  mem[0] & ~mem[1] &  mem[2]  & ~mem[3] &  mem[4]& ~mem[5]

    mem[0] = (mem[0] & ~three135  & ~two14) | two03 | three024
    mem[1] = (mem[1] & ~three024  & ~two25) | two14 | three135
    mem[2] = (mem[2] & ~three135  & ~two03) | two25 | three024
    mem[3] = (mem[3] & ~three024  & ~two14) | two03 | three135
    mem[4] = (mem[4] & ~three135  & ~two25) | two14 | three024
    mem[5] = (mem[5] & ~three024  & ~two03) | two25 | three135

    mem[0] = cyclical_right_shift(mem[0],w)
    mem[1] = cyclical_right_shift(mem[1],w-1)
    mem[2] = cyclical_left_shift(mem[2],1)
    mem[3] = cyclical_left_shift(mem[3],w)
    mem[4] = cyclical_left_shift(mem[4],w-1)
    mem[5] = cyclical_right_shift(mem[5],1)
    
    
def cyclical_right_shift(bs, i):
    mask = 2**i-1
    carry = bs & mask
    slide = bs & ~mask
    return (carry << (l-i)) | (slide>>i)
def cyclical_left_shift(bs, i):
    mask = 2**(l-i) -1
    carry = bs & ~mask
    slide = bs & mask
    return (carry >> (l-i)) | (slide<<i)




def long2bitstring(long):
    bin_str = bin(long)[2:]
    return '0'*(l-len(bin_str)) + bin_str   

def bitCount(int_type):
    count = 0
    while(int_type):
        int_type &= int_type - 1
        count += 1
    return(count)
    
if __name__ == "__main__":
    pygame.init ()
    clock = pygame.time.Clock() 
    
    w = 256
    h = 256
    
    
    screen = pygame.display.set_mode ((2*w + 10, 2*h + 10))
    screen.fill ((250, 250, 250))
    srf = pygame.Surface((w, h))
    scale_srf = pygame.Surface((2*w, 2*h),depth=24,)
    screen.blit (srf, (5, 5))
    l = w*h
    mem = [getrandbits(l),getrandbits(l),getrandbits(l),getrandbits(l),getrandbits(l),getrandbits(l)]
    #mem = [BitVector.BitVector( intVal = 31,size = w*h),BitVector.BitVector(size = w*h),BitVector.BitVector(size = w*h),BitVector.BitVector(intVal = 32,size = w*h),BitVector.BitVector(size = w*h),BitVector.BitVector(size = w*h)]
    #mem = [255,0,0,0,0,255]
    mem =[0,0,0,0,0,0]
     
    for y in xrange(h/4, 3*h/4):
        init_row = 2**(y*w + 3*w/4) - 2**(y*w + w/4)
        mem[0] &= ~init_row
        mem[1] &= ~init_row
        mem[2] &= ~init_row
        mem[3] &= ~init_row
        mem[4] &= ~init_row
        mem[5] &= ~init_row
    for y in xrange(h/4, 2*h/4):
        init_row = 2**(y*w + 2*w/4) - 2**(y*w + w/4)
        mem[0] |= init_row
        mem[1] |= init_row
        mem[2] |= init_row
        mem[3] |= init_row
        mem[4] |= init_row
        mem[5] |= init_row


   


    while True:
        for event in pygame.event.get():
            if event.type == QUIT:
                exit()
            if event.type == KEYUP and event.key == K_ESCAPE:
                exit()
                  
        if False:
            print long2bitstring(mem[0])
            print long2bitstring(mem[1])
            print long2bitstring(mem[2])
            print long2bitstring(mem[3])
            print long2bitstring(mem[4])
            print long2bitstring(mem[5])
            print "-----------"
        
         #print sites
        if False:
            print bmem[0]
            print bmem[1]
            print bmem[2]
            print bmem[3]
            print bmem[4]
            print bmem[5]
            print "-----------"
            
            
            
        #casting longs into bitarrays
        bmem =bitarray()
        for i in xrange(6):
            bmem.extend(long2bitstring(mem[i]))
        
        #casting bitarray bools into np.arrays with true = int(1) and false = int(0)
        bmem = np.fromstring(bmem.unpack(zero='\x00', one='\x01'), dtype=np.uint8)
        #resahpe into 3 dimensions
        bmem = np.reshape(bmem,(6,h,w))
        # summing over the arms of each site to get the mass
        #bmem = np.sum(bmem, axis=0) 
        
        
        srf=pygame.surfarray.make_surface(bmem) 
        srf =   srf.convert(24)
        pygame.transform.scale2x(srf, scale_srf)
        #pygame.transform.smoothscale(srf,(2*w,2*h), scale_srf)
        screen.blit(scale_srf, (5, 5))
        pygame.display.flip()   

       

        step(mem)  
        clock.tick_busy_loop(0)
        framerate = clock.get_fps()
        print framerate