import numpy as np
from matplotlib import animation
from matplotlib import pyplot as plt
import pygame
import sys

#initialize square with values in the middle

def transport(up, down, right, left):
    # is there a way to bit reverse in python?
    new_state = ((right & 0x1) << 3) | ((up & 0x4) >> 1) | ((down & 0x2) << 1) | ((left & 0x8) >> 3)
    return new_state

def run_HPP(NX, NY, square_radius, num_steps):
    all_states = np.zeros((num_steps+1, NX, NY), dtype=np.int)
    # collision rules dict
    col_rules = {0b0000:0b0000, 0b0001:0b1000, 0b0010:0b0100, 0b0011:0b1100, 0b0100:0b0010, 0b0101:0b1010,  
                 0b0110:0b1001, 0b0111:0b1110, 0b1000:0b0001, 0b1001:0b0110, 0b1010:0b0101, 0b1011:0b1101, 
                 0b1100:0b0011, 0b1101:0b1011, 0b1110:0b0111, 0b1111:0b1111}
    
    state = np.zeros((NX, NY), dtype=np.int)
    #random seeding
    for p in np.arange(NX):
      for l in np.arange(NY):
          state[p,l] = np.random.choice(np.arange(16))

    #create center square
    state[NX/2-square_radius:NX/2+square_radius, NY/2-square_radius:NY/2+square_radius] = 0b1111

    new_state = np.zeros((NX, NY), dtype=np.int)
    
    # print "init\n", state
    all_states[0,:,:] = state
    for step in np.arange(num_steps):
        # apply collision rules
        for i in np.arange(NX):
            for j in np.arange(NY):            
                new_state[i,j] = col_rules[state[i,j]]

        print "post collision\n", new_state
        # apply transport rules
        for i in np.arange(1, NX-1):
            for j in np.arange(1, NY-1):            
                state[i,j] = transport(new_state[i, j+1], new_state[i, j-1], new_state[i+1, j], new_state[i-1, j])

        print "post transport\n", state

        #apply periodic boundaries
        for i in np.arange(NX):
            state[i,0] = state[0, NX-2]
            state[i,NX-1] = state[i,1]

        for j in np.arange(NY):
            state[0,j] = state[NX-2,j]
            state[NX-1,j] = state[1,j]
            
        all_states[step+1,:,:] = state
            
    return all_states

###########################################
# NOTE: Animation Modeled after Sam's Code
###########################################
def draw_states(state, NX, NY):
    screen.fill((200,200,220))
    pxls = 30
    for i in range(NX):
        for j in range(NY):
            pygame.draw.circle(screen, (0,255,0), (int(pxls*i),int(pxls*j)), int(.2*pxls))
            if state[i,j] & 0x1: pygame.draw.circle(screen, (240,0,0), (int(pxls*(i-.25)),int(pxls*j)), int(.1*pxls))
            if state[i,j] & 0x2: pygame.draw.circle(screen, (240,0,0), (int(pxls*i),int(pxls*(j-.25))), int(.1*pxls))
            if state[i,j] & 0x4: pygame.draw.circle(screen, (240,0,0), (int(pxls*i),int(pxls*(j+.25))), int(.1*pxls))
            if state[i,j] & 0x8: pygame.draw.circle(screen, (240,0,0), (int(pxls*(i+.25)),int(pxls*j)), int(.1*pxls))
    pygame.display.update()

if __name__ == "__main__":
    NX = 20
    NY = 20
    square_radius = 5
    num_steps = 30
    M = run_HPP(NX, NY, square_radius, num_steps)
    print M.shape

    pygame.init()
    pxls = 30
    screen = pygame.display.set_mode((pxls*NX,pxls*NY))
    clock = pygame.time.Clock()

    for step in np.arange(num_steps):
        msElapsed = clock.tick(5)
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit(); sys.exit();

        draw_states(M[step], NX, NY)
        pygame.image.save(screen,"hpp_random/%04d.png"%step)