from matrix import get_size, pmx
import imageio
import matplotlib.pyplot as plt
import numpy as np


class Animate:
    def __init__(self, mx, day="CHANGEME"):
        self.mx = mx
        self.day = day
        _size = get_size(mx)
        self.height = _size[0]
        self.width = _size[1]
        self.f_count = -1

    def add_frame(self, frame):
        self.f_count += 1
        self.write_frame(frame)

    def write_frame(self, frame):
        current = np.zeros_like(self.mx)
        for y, row in enumerate(frame):
            for x, col in enumerate(row):
                current[y][x] = frame[y][x]
        fig, ax = plt.subplots()
        ax.imshow(current, cmap=plt.cm.gray)
        ax.axis("off")
        _figpath = f"gif-{self.day}/{self.f_count:05}.png"
        plt.savefig(_figpath)
        plt.close()

    def animate(self, frameskip=1):
        with imageio.get_writer(
            f"gif-{self.day}/day{self.day}.gif", mode="I"
        ) as writer:
            names = [
                f"gif-{self.day}/{x:05}.png" for x in range(0, self.f_count, frameskip)
            ]
            print(names)
            for filename in names:
                try:
                    image = imageio.imread(filename)
                    writer.append_data(image)
                except FileNotFoundError:
                    pass