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(east, northeast, southeast, west, northwest, southwest):
    # new_state = ((east & 0b100000) >> 5) | ((northeast & 0b010000) >> 3) | ((southeast & 0b001000) >> 1) | ((west & 0b000001) << 5) | ((northwest & 0b000100) << 1) | ((southwest & 0b000010) << 3)
    new_state = (east & 0b100000) | (northeast & 0b010000) | (southeast & 0b001000) | (west & 0b000001) | (northwest & 0b000100) | (southwest & 0b000010)
    # bit reverse!
    new_state = int(bin(new_state)[:1:-1], 2)
    return new_state

def run_FHP(NX, NY, square_radius, num_steps):
    all_states = np.zeros((num_steps+1, NX, NY), dtype=np.int)
    # collision rules dict
    col_rules_one = {0b001100:0b100001, 0b010010:0b100001, 0b100001:0b010010, 0b011001:0b011001, 0b100110:0b100110}
    col_rules_two = {0b001100:0b010010, 0b010010:0b001100, 0b100001:0b001100, 0b011001:0b011001, 0b100110:0b100110}
    
    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(64))

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

    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):
                if state[i,j] in col_rules_one.keys():
                    new_state[i,j] = col_rules_one[state[i,j]]
                else: # bit reverse
                    new_state[i,j] = int(bin(state[i,j])[:1:-1], 2)

        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-1, j], new_state[i-1, j+1], new_state[i-1, j-1], new_state[i+1, j], new_state[i+1, j+1], new_state[i+1, j-1])

        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
    dot = 0.35
    for i in range(NX):
        for j in range(NY):
            dx = (j%2)/2.
            pygame.draw.circle(screen, (0,255,0), (int(pxls*(i+dx)),int(pxls*j)), int(.2*pxls))
            if state[i,j] & 1: pygame.draw.circle(screen, (240,0,0), (int(pxls*(i-dot+dx)),int(pxls*j)), int(.1*pxls))
            if state[i,j] & 2: pygame.draw.circle(screen, (240,0,0), (int(pxls*(i-.5*dot+dx)),int(pxls*(j+.866*dot))), int(.1*pxls))
            if state[i,j] & 4: pygame.draw.circle(screen, (240,0,0), (int(pxls*(i-.5*dot+dx)),int(pxls*(j-.866*dot))), int(.1*pxls))
            if state[i,j] & 8: pygame.draw.circle(screen, (240,0,0), (int(pxls*(i+.5*dot+dx)),int(pxls*(j+.866*dot))), int(.1*pxls))
            if state[i,j] & 16: pygame.draw.circle(screen, (240,0,0), (int(pxls*(i+.5*dot+dx)),int(pxls*(j-.866*dot))), int(.1*pxls))
            if state[i,j] & 32: pygame.draw.circle(screen, (240,0,0), (int(pxls*(i+dot+dx)),int(pxls*j)), int(.1*pxls))
    pygame.display.update()

if __name__ == "__main__":
    NX = 20
    NY = 20
    square_radius = 5
    num_steps = 30
    M = run_FHP(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(10)
        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,"fhp_frames_random/%04d.png"%step)