import matplotlib.animation as animation
import matplotlib.pyplot as plt
#import plotly.graph_objects as go
import plotly.express as px
import numpy as np


class PlotlyAnimation(object):
    def __init__(self, state, steps):
        self.state = state
        self.steps = steps
        self.frames = np.expand_dims(self.state.board, axis=0)
        self.frame_gen()
        self.animation = self.animation_gen()

    def frame_gen(self):
        for _ in range(self.steps):
            self.frames = np.concatenate((self.frames, np.expand_dims(self.state.step(), axis=0)), axis=0)
        return self.frames

    def animation_gen(self):
        fig = px.imshow(self.frames, animation_frame=0)
        fig.layout.updatemenus[0].buttons[0].args[1]['frame']['duration'] = 30
        fig.layout.updatemenus[0].buttons[0].args[1]['transition']['duration'] = 5
        fig.layout.updatemenus[0].showactive = True
        fig.layout.sliders[0].visible = False
        fig.layout.coloraxis.showscale = False
        fig.layout.xaxis.showticklabels = False
        fig.layout.yaxis.showticklabels = False
        return fig

    def show(self):
        self.animation.show()


class PltAnimation(object):
    def __init__(self, state, steps=None):

        self.state = state
        self.steps = steps
        self.fig, self.ax = plt.subplots()
        self.ax = self.ax.matshow(self.state.board)

    def frame_yield(self):
        while True:
            self.state.step()
            yield self.state.board

    def frame_gen(self):
        frames = []
        for i in range(self.steps):
            self.state.step()
            [].append(self.state.board)
        return self.ax, frames

    def init(self):
        return self.state.board

    def update(self, data):
        self.ax.set_data(data)
        return self.ax

    def animate(self):
        if self.steps is None:
            return animation.FuncAnimation(
                self.fig,
                self.update,
                self.frame_yield,
                init_func=self.init,
                interval=100,
            )
        else:
            return animation.FuncAnimation(
                self.fig,
                self.update,
                self.frame_gen,
                init_func=self.init,
                interval=100,
            )