import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.colors import LogNorm

EVAPORATE = 0.01
DIFFUSION = 1
INITDROP = 20
LOWERBOUND = 0.01

def von_neuman_neighbors(i, j, size):
    n = []
    if i > 0:
        n.append((i - 1, j))
    if j > 0:
        n.append((i, j - 1))
    if i < size - 1:
        n.append((i + 1, j))
    if j < size - 1:
        n.append((i, j + 1))
    return n

def moore_neighbors(i, j, size):
    n = []
    if i > 0:
        n.append((i - 1, j))
        if j > 0:
            n.append((i - 1, j - 1))
        if j < size - 1:
            n.append((i - 1, j + 1))
    if j > 0:
        n.append((i, j - 1))
    if i < size - 1:
        n.append((i + 1, j))
        if j > 0:
            n.append((i + 1, j - 1))
        if j < size - 1:
            n.append((i + 1, j + 1))
    if j < size - 1:
        n.append((i, j + 1))
    return n


class CA():
    
    def __init__(self, size):
        self.size = size
        self.x = np.arange(size)
        self.y = np.arange(size)
        self.cells = np.zeros((size, size))
        
    def pattern_setup(self, pattern):
        for p in pattern:
            self.cells[p[0] + self.size / 2, p[1] + self.size / 2] = INITDROP
        
    def random_setup(self, n):
        for i in range(n):
            a, b = np.random.random_integers(0, size - 1, 2)
            self.cells[a, b] = INITDROP
        
    def image_setup(self):
        self.plt = plt.imshow(self.cells, interpolation='nearest', 
                            origin='bottom', 
                            vmin=np.min(LOWERBOUND),
                            vmax=np.max(INITDROP),
                            norm=LogNorm(vmin=LOWERBOUND, vmax=INITDROP),
                            cmap=plt.cm.YlGn)

    def update(self):
        newcells = np.zeros((self.size, self.size))
        for i in range(self.size):
            for j in range(self.size):
                neighbors = moore_neighbors(i, j, self.size)               
                otherP = self.cells[i,j]
                for a, b in neighbors:
                    otherP += self.cells[a,b]
                otherP /= len(neighbors) + 1
                newcells[i,j] = (1 - EVAPORATE) * (self.cells[i,j] + (DIFFUSION * (otherP - self.cells[i,j])))
                if newcells[i,j] < LOWERBOUND:
                    newcells[i,j] = 0
        self.cells = newcells
        
    def plot(self):
        self.plt.set_data(self.cells)
        return self.plt
        
size = 81

fig, ax = plt.subplots()
ax.set_ylim(-1, size)
ax.set_xlim(-1, size)

factories = ((1, 2), (10, 1), (-7, -10), (15, 2), (30, -25))
       
ca = CA(size)
ca.pattern_setup(factories)
ca.image_setup()

def update(data):
    ca.update()
    return ca.plot(),

def data_gen():
    while True: yield 1

ani = animation.FuncAnimation(fig, update, data_gen, blit=False, interval=100)
plt.show()