import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation
from tensorflow.keras.layers import Conv2D, InputLayer
from tensorflow.keras.models import Sequential

size = (10, 10)
expanded_size = (1, *size, 1)
n_frames = 120
glider = ((1, 2), (2, 3), (3, 1), (3, 2), (3, 3))

# env = np.random.randint(0, 2, expanded_size)
env = np.zeros(expanded_size, dtype=int)
for pos in glider:
    env[(0,) + pos] = 1

model = Sequential([InputLayer(input_shape=expanded_size[1:]),
                    Conv2D(1, 3, padding="same", activation=None, use_bias=False, kernel_initializer="ones")])

frames = []
for i in range(n_frames):
    # 2D sliding window of 3x3 including summation TODO: add torus-style wrapping
    # NOTE: convolve2d of scipy does support torus-padding but that's obviously not as cool as a neural network
    neighbours = model(env) - env  # Don't count the cell itself in the number of neighbours
    env = np.where((env & np.isin(neighbours, (2, 3))) | ((env == 0) & (neighbours == 3)), 1, 0)
    frames.append(env.squeeze())


fig = plt.figure()
ax = plt.axes(xlim=(0, size[0]), ylim=(0, size[1]))
render = plt.imshow(frames[0], interpolation="none", cmap="binary")


def animate(i: int):
    render.set_array(frames[i])
    return [render]


anim = FuncAnimation(fig, animate, frames=n_frames, interval=30, blit=True)
plt.axis("off")
plt.gca().invert_yaxis()
anim.save("glider.gif", fps=30)
plt.show()