# neil 2018
# illustration of Conway's Game of Life
# this version uses 'imshow' (which tends to slow down because imshow stores 
# all the figure data, which can get very large)
# See life_animated.py for a better visualization method

# heavily modified from a program on GitHub
# The animation routines may not work unless the graphics
# is set to qt5 or 'automatic', I don't think it works 'inline'

import numpy as np
import matplotlib.pyplot as plt


seeds = {'diehard': [[0, 0, 0, 0, 0, 0, 1, 0],
                     [1, 1, 0, 0, 0, 0, 0, 0],
                     [0, 1, 0, 0, 0, 1, 1, 1]],
         'boat': [[1, 1, 0],
                  [1, 0, 1],
                  [0, 1, 0]],
         'r_pentomino': [[0, 1, 1],
                         [1, 1, 0],
                         [0, 1, 0]],
         'pentadecathlon': [[1, 1, 1, 1, 1, 1, 1, 1],
                            [1, 0, 1, 1, 1, 1, 0, 1],
                            [1, 1, 1, 1, 1, 1, 1, 1]],
         'beacon': [[1, 1, 0, 0],
                    [1, 1, 0, 0],
                    [0, 0, 1, 1],
                    [0, 0, 1, 1]],
         'acorn': [[0, 1, 0, 0, 0, 0, 0],
                   [0, 0, 0, 1, 0, 0, 0],
                   [1, 1, 0, 0, 1, 1, 1]],
         'spaceship': [[0, 0, 1, 1, 0],
                       [1, 1, 0, 1, 1],
                       [1, 1, 1, 1, 0],
                       [0, 1, 1, 0, 0]],
         'block_switch_engine': [[0, 0, 0, 0, 0, 0, 1, 0],
                                 [0, 0, 0, 0, 1, 0, 1, 1],
                                 [0, 0, 0, 0, 1, 0, 1, 0],
                                 [0, 0, 0, 0, 1, 0, 0, 0],
                                 [0, 0, 1, 0, 0, 1, 0, 0],
                                 [1, 0, 1, 0, 0, 0, 0, 0]],
         'infinite': [[1, 1, 1, 0, 1],
                      [1, 0, 0, 0, 0],
                      [0, 0, 0, 1, 1],
                      [0, 1, 1, 0, 1],
                      [1, 0, 1, 0, 1]]}


def survival(x, y, universe):
    """
    Calculates if a cell at x,y lives (1) or dies (0)
    'universe' is the full array of zeros and ones 
    """
    num_neighbours = np.sum(universe[x-1:x+2, y-1:y+2]) - universe[x, y]

    # The rules of Life
    s = universe[x,y]     # 's' is either 0 or 1                
    if s == 1 and num_neighbours < 2 or num_neighbours > 3:
        s = 0             # the cell dies if less than 2, more than 3 neighbors
    elif s == 0 and num_neighbours == 3:
        s = 1             # cell comes alive if dead and exactly 3 neighbours
    return s              # otherwise live cells stay alive, and dead ones dead.

def generation(universe):
    """
    ranges over the entire array, and populates a new version of the universe
    returns the new universe
    """
    new_universe = np.copy(universe)
    # Apply the survival function to every cell in the universe
    for i in range(universe.shape[0]):
        for j in range(universe.shape[1]):
            new_universe[i, j] = survival(i, j, universe)
    return new_universe

# Start of main program, which will call the 2 above functions
# header data
universe_size=(100, 100) #size of universe, note this does not have periodic BCs
interval=100             #time in millesecs between display steps
n_generations=500        # number of 'time' steps

# the following 6 lines allow loading one of the above 'seed' patterns
seed='spaceship'
seed_position=(int(universe_size[0]*4/10), int(universe_size[1]*4/10))
# Initialise the universe and seed
universe = np.zeros(universe_size)
x_start, y_start = seed_position[0], seed_position[1]
seed_array = np.array(seeds[seed])
universe[x_start:x_start+seed_array.shape[0], y_start:y_start+seed_array.shape[1]] = seed_array
# instead, you can start with a random array if you un-delete the following line
#universe = np.random.randint(0,2,universe_size)

# Animate
fig = plt.figure()                    #create a place to plot
ax  = fig.add_subplot(111)            #create axes.  you can create multiple plots on the same figure with subplot(2,3,1) example


for i in range(n_generations):
    print(i)                          # so we can see where we are
    im=ax.imshow(universe, cmap='Purples')
    universe = generation(universe)
    plt.pause(.1)                     # if this is not included the plots don't appear on the screen

plt.show()