import matplotlib.animation as animation
import matplotlib.pyplot as plt
import cellpylib as cpl

class CellularAutomata:
  def __init__(self, size:int=10, timesteps:int=50):
    self.size = size
    self.timesteps = timesteps
    self.cellular_automata = cpl.init_random2d(size, size)

  def set_game_of_live_roles(self):
    self.cellular_automata = cpl.evolve2d(self.cellular_automata, timesteps=self.timesteps, neighbourhood='Moore', apply_rule=cpl.game_of_life_rule, memoize='recursive')

  def get_ready_for_animation(self):
    fig, ax = plt.subplots()
    ax.set_xlim((0, self.size))
    ax.set_ylim((0, self.size))

    self.img = ax.imshow(self.cellular_automata[0], interpolation="nearest", cmap='Greys' )
    return fig

  def init_animation(self):
    self.img.set_data(self.cellular_automata[0])
    return (self.img, )

  def make_frame(self, i):
    self.img.set_data(self.cellular_automata[i])
    return (self.img, )

  def run_animation(self):
    fig = self.get_ready_for_animation()
    ani = animation.FuncAnimation(fig, self.make_frame, init_func=self.init_animation, frames=self.timesteps, interval=30, blit=True, repeat=False)
    plt.show()

if __name__ == '__main__':
  cellular_automata = CellularAutomata()
  cellular_automata.set_game_of_live_roles()
  cellular_automata.run_animation()