CA_neighbour.py

##### RANDOM MAP GENERATOR #####
import pygame
import random
import math
import time


class MapGrid():
    def __init__(self, map_width, map_height,moi):

        # set map values
        self.map_width = map_width
        self.map_height = map_width
        self.moi = moi

        # generate outside rooms
        self.outside_terrain_grid = self._generate_init_virus_grid(self.map_width, self.map_height,self.moi)

    def _index_transfer(self, map_width, map_height, one_index):
        '''transfer 1d index (one_index) to 2d'''
        x = one_index % map_width
        y = one_index / map_width
        return [x,y]

    def _generate_init_virus_grid(self, map_width, map_height, moi):
        '''
        creates a random distirubted virus in the map with moi from 0 to 1
        '''
        new_map_grid = [[0]* map_width for i in xrange(map_height)] # create our new list
        num_grids = map_height*map_width
        virus_seq = random.sample(range(num_grids),int(moi*num_grids))
        
        # get the each virus index in 2d form
        for virus in virus_seq:
            virus_ind = self._index_transfer(map_width,map_height,virus)
            new_map_grid[virus_ind[0]][virus_ind[1]] = 1

        return new_map_grid


    def _generate_outside_terrain(self, empty_outside_terrain_grid, number_of_generations):
        '''
        creates a bubble effect with cellular automaton
        '''
        grid = empty_outside_terrain_grid
        number_of_generations = number_of_generations

        for x in range(number_of_generations):
            next_grid = []
            for column_index, column in enumerate(grid):
                next_column = []
                next_grid.append(next_column)
                for tile_index, tile in enumerate(column):

                    # get the surrounding tile values for each tile
                    top_left = grid[column_index - 1][tile_index - 1]
                    top_mid = grid[column_index][tile_index - 1]
                    try:
                        top_right = grid[column_index + 1][tile_index - 1]
                    except IndexError:
                        top_right = 0
                    
                    center_left = grid[column_index - 1][tile_index]
                    center_mid = grid[column_index][tile_index]
                    try:
                        center_right = grid[column_index + 1][tile_index]
                    except IndexError:
                        center_right = 0

                    try:
                        bottom_left = grid[column_index - 1][tile_index + 1]
                    except IndexError:
                        bottom_left = 0
                    try:
                        bottom_mid = grid[column_index][tile_index + 1]
                    except IndexError:
                        bottom_mid = 0
                    try:
                        bottom_right = grid[column_index + 1][tile_index + 1]
                    except IndexError:
                        bottom_right = 0


                    neighbors = [top_mid, center_left , center_mid, center_right ,
                                 bottom_mid, top_left , top_right , bottom_left , bottom_right]

                    number_of_A1 = neighbors.count(1)
                    number_of_A2 = neighbors.count(2)

                    if tile ==0: #healthy
                        if number_of_A1 > 0:
                            next_cell = 1
                        elif number_of_A2 >=4:
                            next_cell = 1
                        else:
                            next_cell = tile
                    elif tile == 1:
                        next_cell = 2 
                    elif tile ==2:
                        next_cell = 3
                    elif tile == 3:#dead
                        if random.random()<.99:
                            if random.random() < .00001:
                                next_cell = 1 
                            else:
                                next_cell = 0 #replace by healthy
                        else:
                            next_cell = tile # remain dead    


                    # create the new cell
                    next_column.append(next_cell)
            grid = next_grid
        return next_grid

if __name__ == '__main__':
    
    # general map stats
    map_width = 100
    map_height = 100
    moi = 0.1
    # start with one generation
    tile_size = 8
    unit_size = 8
        
    map_grid = MapGrid(map_width, map_height,moi) #init a random grid filled with 0s and 1s 
    print sum(sum(x) for x in map_grid.outside_terrain_grid[:])

    pygame.init()

    screen = pygame.display.set_mode((map_width * tile_size,map_height * tile_size))#return Surface

    zero_tile = pygame.Surface((unit_size, unit_size))
    zero_tile.fill((0,0,255))#blue
    one_tile = pygame.Surface((unit_size,unit_size))
    one_tile.fill((255,205,0)) # yellow 
    two_tile = pygame.Surface((unit_size,unit_size))
    two_tile.fill((0,128,0)) # green
    three_tile = pygame.Surface((unit_size,unit_size))
    three_tile.fill((196,2,51)) # red
    colors = {0: zero_tile, 1: one_tile,2: two_tile , 3: three_tile} # 0: healthy-blue, 1: infected-A1-yellow, 2:

    background = pygame.Surface((map_width * tile_size,map_height * tile_size))

    clock = pygame.time.Clock()#an object to help track time 

    first_gen = True
    timer = 1
    
    running = True
    while running == True:
        clock.tick(1)
        pygame.display.set_caption('Infection lasts: ' + str(pygame.time.get_ticks())+' weeks')
        for event in pygame.event.get():#get the events from queue
            if event.type == pygame.QUIT:
                running = False

        if first_gen:
            themap = map_grid.outside_terrain_grid
        else:
            themap = map_grid._generate_outside_terrain(themap, 1) #update 1 generation 

        for column_index, column in enumerate(themap):
            for tile_index, tile in enumerate(column):
                screen.blit(colors[tile], (tile_index * tile_size, column_index * tile_size))

        pygame.display.flip()

        if first_gen:
            timer -= 1
            if timer < 0:
                first_gen = False
        print timer
    pygame.quit()