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


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

    def add_frame(self, frame):
        self.frames.append(deepcopy(frame))

    def print(self):
        for f in self.frames:
            pmx(f)
            print("-")


    def animate(self, save=False, or_frames=False):
        last_file = 0
        for f_idx, frame in enumerate(self.frames):
            current = np.zeros_like(self.mx)
            for y, row in enumerate(frame):
                for x, col in enumerate(row):
                    current[y][x] = frame[y][x]
            pmx(current)
    
            fig, ax = plt.subplots()
            ax.imshow(current, cmap=plt.cm.gray)
            ax.axis("off")
            if save:
                plt.savefig(f"{self.day}/{f_idx:03}.png")
            plt.close()
    
        with imageio.get_writer(f"{self.day}/day{self.day}.gif", mode="I") as writer:
            names = [f"{self.day}/{x:03}.png" for x in range(len(self.frames))]
            for filename in names:
                try:
                    image = imageio.imread(filename)
                    writer.append_data(image)
                except FileNotFoundError:
                    pass