from pylab import*
from numpy import*
import matplotlib.pyplot as plt
from matplotlib import animation, rc
from IPython.display import HTML
from time import sleep
import sys
sys.setrecursionlimit(200000)

def povecaj(i,j):
    global event
    rob=False#True
    #random.randint(8, size=1)[0]
    ip=int(i)+1
    im=int(i)-1
    jp=int(j)+1
    jm=int(j)-1
    '''if im>0 and ip<(N)and jm>0 and jp<(N):
        rob=False'''
    #print i,j,rob,S[i,j]>3
    if rob==False and S[i,j]>3.5:
        S[i,j]-=4
        if ip<N:
            S[ip,j]+=1
            povecaj(ip,j)
            #print "ip",ip,S[ip,j]
        if im>0:
            S[im,j]+=1
            povecaj(im,j)
            #print "im",im,S[im,j]
        if jp<N:
            S[i,jp]+=1
            povecaj(i,jp)
            #print "jp",jp,S[i,jp]
        if jm>0:
            S[i,jm]+=1
            povecaj(i,jm)
            #print "jm",jm,S[i,jm]
        #print "--------------"
        

N=21
S=zeros([N,N],int)

fig,ax =plt.subplots()
surf=ax.imshow(S,vmin=0, vmax=8, alpha=1,interpolation="none", cmap="hot")
surf.set_data(S)

freq={}
plaz=[]

def animate(i):
    global freq
    global plaz
    ii=int((N-1)/2)
    jj=int((N-1)/2)
    event=0
    S[int(ii),int(jj)]+=1
    povecaj(ii,jj)
    if event!=0:
        plaz.append(event)
    if event in freq:
        freq[event]+=1
    else:
        freq[event]=1
    surf.set_data(S)
    return(surf,)

anim = animation.FuncAnimation(fig, animate,frames=1000, interval=25, blit=True)
Writer = animation.writers['ffmpeg']
writer = Writer(fps=25, metadata=dict(artist='Rene_Markovic'), bitrate=1800)
#anim.save('TBW_model.mp4', writer=writer)
HTML(anim.to_html5_video())