import random
import copy
import matplotlib.pyplot as plt
import matplotlib.animation as animation


class State(object):
    def __init__(self, depth: int, width: int, coords=None):
        # Coords: tuple of ints or list of tuple of ints
        self.width = width
        self.depth = depth

        self.board = [[0 for i in range(width)] for j in range(depth)]
        if coords is None:
            pass
        elif isinstance(coords, tuple):
            x, y = coords
            self.board[x][y] = 1
        elif isinstance(coords, list):
            for pair in coords:
                x, y = pair
                self.board[x][y] = 1
        elif coords == "random":
            self.randomBoard()
        else:
            raise Exception(
                "No proper coords provided. Try a tuple of ints or a list of tuple of ints"
            )

    def randomBoard(self):
        for idx, i in enumerate(self.board):
            for jdx, j in enumerate(i):
                if random.randint(1, 2) == 1:
                    self.board[idx][jdx] = 1
                else:
                    self.board[idx][jdx] = 0

    def terminal_render(self):
        print("-" * (self.width + 2))
        for i in self.board:
            print(
                "|"
                + "".join(str(k) for k in ([" " if j == 0 else "*" for j in i]))
                + "|"
            )
        print("-" * (self.width + 2))

    def plt_render(self):
        plt.imshow(self.board)
        plt.show()

    def step(self):
        boardcopy = copy.deepcopy(self.board)
        for rowidx, row in enumerate(self.board):
            for elemidx, elem in enumerate(row):
                live_neighbours = self.checkNeighbours(boardcopy, rowidx, elemidx)
                if self.board[rowidx][elemidx] == 1:
                    if live_neighbours == 2 or live_neighbours == 3:
                        pass
                    else:
                        self.board[rowidx][elemidx] = 0
                else:
                    if live_neighbours == 3:
                        self.board[rowidx][elemidx] = 1

    def checkNeighbours(self, board, rowindex, columnindex, neighbourhood="moore"):
        alive = 0
        if rowindex != 0:
            if board[rowindex - 1][columnindex] == 1:
                alive += 1
        if rowindex < self.depth - 1:
            if board[rowindex + 1][columnindex] == 1:
                alive += 1
        if columnindex != 0:
            if board[rowindex][columnindex - 1] == 1:
                alive += 1
        if columnindex < self.width - 1:
            if board[rowindex][columnindex + 1] == 1:
                alive += 1
        if neighbourhood == "moore":
            if rowindex != 0 and columnindex != 0:
                if board[rowindex - 1][columnindex - 1] == 1:
                    alive += 1
            if rowindex < self.depth - 1 and columnindex < self.width - 1:
                if board[rowindex + 1][columnindex + 1] == 1:
                    alive += 1
            if rowindex != 0 and columnindex < self.width - 1:
                if board[rowindex - 1][columnindex + 1] == 1:
                    alive += 1
            if rowindex < self.depth - 1 and columnindex != 0:
                if board[rowindex + 1][columnindex - 1] == 1:
                    alive += 1
        return alive


class BoardAnimation(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):
        return animation.FuncAnimation(
            self.fig, self.update, self.frame_yield, init_func=self.init, interval=100
        )


def main():
    state = State(20, 20, [(0, 1), (1, 2), (2, 0), (2, 1), (2, 2)])
    animation = BoardAnimation(state)
    ani = animation.animate()
    plt.show()


if __name__ == "__main__":
    main()