import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation
from functools import partial

plt.style.use("ggplot")

#initializes static elements
def init_fig(fig, axes, artist):
    axes.axes.get_xaxis().set_visible(False)
    axes.axes.get_yaxis().set_visible(False)
    axes.get_xaxis().tick_bottom()
    artist.set_clim([0.0, 1.0])
    artist.set_cmap("binary")
    artist.set_interpolation("quadric")
    fig.colorbar(heat_map)
    return artist,

#returns data needed to update artist
def animate(t, N, heat_map, A, B, D_A, D_B, conv, f, k):
    print(t)

    #derivative matrix
    dA = np.zeros([N, N])
    dB = np.zeros([N, N])

    #iterate over each non-border cell
    for i in range(1, N-1):
        for j in range(1, N-1):
            dA[i, j] = D_A*laplace(A, (i, j), conv) - A[i, j]*(B[i, j])**2 + f*(1-A[i, j])
            dB[i, j] = D_B*laplace(B, (i, j), conv) + A[i, j]*(B[i, j])**2 - (k+f)*B[i, j]

    #evolving
    A += dA
    B += dB

    heat_map.set_data(A)
    return heat_map,

#Laplace operator at a point
def laplace(matrix, cell, conv):
    values = matrix[cell[0]-1 : cell[0]+2, cell[1]-1 : cell[1]+2]
    return np.sum(conv*values)

#initial data
conv = np.array([[0.05, 0.2, 0.05], [0.2, -1.0, 0.2], [0.05, 0.2, 0.05]])
D_A, D_B, f, k = 1.0, 0.5, 0.0545, 0.062
t = 100

N = 100
A, B = np.ones([N, N]), np.zeros([N, N])
for i in range(-int(N/4), int(N/4)):
    for j in range(-int(N/4), int(N/4)):
        B[int(N/2)-10+i, int(N/2)-10+j] = 1.0

#figure setup
fig, axes = plt.figure(), plt.axes(frameon = True)
heat_map = axes.imshow(A, animated = True)

init = partial(init_fig, fig=fig, axes=axes, artist=heat_map)

#animate
anim = FuncAnimation(fig, func=animate, frames=t, init_func=init, blit=True, fargs=(N, heat_map, A, B, D_A, D_B, conv, f, k))
anim.save("out.mp4", fps=40)