import numpy as np
import seaborn as sns; sns.set()
%matplotlib inline

# Steps:
# 1. Generate a large matrix - fill with zeroes
# 2. Initialize some structure as a starting point.
# 3. Calculate elements that are dying.
# 4. Calculate elements that are being born.
# 5. Add new elements and remove the died ones.
# 6. plot....


def get_dying(matrix):
    '''
    This function returns a list of coordinates of the dying
    cells.
    '''
    nonZero = np.nonzero(matrix)
    dying = []
    for index, x in enumerate(nonZero[0]):
        # Find dying cells:
        y = nonZero[1][index]
        # Get all non-zero element.

        # Dying because of of overpopulation (More than 3 neighbour):
        if  get_neighbours(x,y,matrix) > 3:
            dying.append((x,y))

        # Dying because of of under population (Less than two neighbour):
        if get_neighbours(x,y,matrix) < 2:
            dying.append((x,y))

    # Returning the coordinates of dying cells:
    return dying

def get_neighbours(x,y,matrix):
    '''
    given the x and y coordnates of a matrix, this function
    returns the number of non-zero neighbours of given coordinate
    '''
    # print matrix[(x-1):(x+2), (y-1):(y+2)]
    if matrix[x,y] == 1:
        return  np.count_nonzero(matrix[(x-1):(x+2), (y-1):(y+2)]) - 1
    else:
        return np.count_nonzero(matrix[(x-1):(x+2), (y-1):(y+2)])

def get_born(matrix):
    '''
    this function returns with a list of tuples containing the
    coordinates of newly borning cells.
    '''

    nonZero = np.nonzero(matrix)
    born = []
    for index, x in enumerate(nonZero[0]):
        # Find dying cells:
        y = nonZero[1][index]

        # Checking zero elements around the non-zero elements:
        for dx in [x - 1, x, x + 1]:
            for dy in [y - 1, y, y + 1]:
                if 0 <= dx < width and 0 <= dy < height:
                    if get_neighbours(dx, dy, matrix) == 3:
                        born.append((dx,dy))
    return born


# initialize starting structure:
def initialize(matrix):
    '''
    This function takes a matrix as an input and at a
    random position creates a five element big "patch of ones"
    '''

    # get dimension of the matrix:
    (w, h) = np.shape(matrix)

    # Get random coordinate:
    x = np.random.choice(w)
    y = np.random.choice(h)

    # Update matrix:
    matrix[x,y] = 1

    for i in range(8):
        # Updating x and y:
        (dx, dy) = np.random.choice([-1,0,1], 2)
        if 0 < x + dx < w and 0 < y + dy < h:
            i += 1
            x = x + dx
            y = y + dy
            matrix[x,y] = 1

    # Returning the modified matrix
    return matrix



def get_new_gen(matrix):

    # Before any modification is done, we get the
    # list of cells:
    dying_cells = get_dying(matrix)
    borning_cells = get_born(matrix)

    # Updating matrix:
    for dying in dying_cells:
        # print "dying: ", dying[0], dying[1]
        matrix[dying[0], dying[1]] = 0

    for born in borning_cells:
        # print "born: ", born[0], born[1]
        matrix[born[0], born[1]] = 1

    return matrix


# Initialize matrix with a given size:
width = 200
height = 200
matrix = np.zeros((width, height))

## Initialize 20 starting object:
for i in range(60):
    matrix = initialize(matrix)

# Plot the initial structure:
sns.heatmap(matrix, xticklabels=False, yticklabels=False, cbar=False)

# Create generations (Let's say 250):
for i in range(250):
    get_new_gen(matrix)
    ## Plotting heatmap:
    plot = sns.heatmap(matrix, xticklabels=False, yticklabels=False, cbar=False)
    x = plot.get_figure()
    filename = "output_%02d.png" % i
    x.savefig(filename)