import numpy as np
from scipy import signal
from PIL import Image

from utils import generate_image_from_grid


class GameOfLife:

    def __init__(self):
        """Initialize the game.
        dim: dimensions of the board
        p: probability of a cell being dead at init
        seed: (optional) for reproducibility, set to None for a new random state
        init_state: (optional) a np.array grid to start the game with
        """
        self.kernel = [
            [1, 1, 1],
            [1, 0, 1],
            [1, 1, 1],
        ]

        self.kernel = np.ones((3, 3))
        self.kernel[1, 1] = 0
        self.game_history = []
        self.state = None
        self.step_counter = 0

    def set_random_state(self, dim=(100, 100), p=0.5, seed=None):
        if seed:
            np.random.seed(seed)
        self.state = (np.random.random(dim) < p).astype("int")
        self.game_history.append(self.state.copy())

    def set_empty_state(self, dim=(100, 100)):
        self.state = np.zeros(dim)
        self.game_history.append(self.state.copy())

    def set_state_from_array(self, array):
        self.state = array.copy()
        self.game_history.append(self.state.copy())

    def count_neighbors(self):
        """
        Count the number of live neighbors each cell in self.state has with convolutions.
        """
        self.neighbors = signal.convolve2d(
            self.state, self.kernel, boundary="fill", fillvalue=0, mode="same"
        ).astype("int")

    def place_blob(self, blob, i, j):
        """Place a blob at coordinates i,j
        blob: ndarray of zeros and ones
        i: int
        j: int
        """
        try:
            self.state[i : i + blob.shape[0], j : j + blob.shape[1]] = blob
        except:
            print("Check bounds of box vs size of game!")

    def step(self):
        """Update the game based on conway game of life rules"""
        # Count the number of neighbors via convolution
        self.count_neighbors()

        # Copy of initial state
        self.new_state = self.state

        # Rebirth if cell is dead and has three live neighbors
        self.new_state += np.logical_and(self.neighbors == 3, self.state == 0)

        # Death if cell has less than 2 neighbors
        self.new_state -= np.logical_and(self.neighbors < 2, self.state == 1)

        # Death if cell has more than 3 neighbors
        self.new_state -= np.logical_and(self.neighbors > 3, self.state == 1)

        # Update game state
        self.state = self.new_state

        # Save state to history
        self.game_history.append(self.state.copy())

        # Update step counter
        self.step_counter += 1

    def generate_n_steps(self, n):
        for _ in range(n):
            self.step()

            if np.array_equal(self.game_history[-1], self.game_history[-2]):
                # If the game is stable, break
                break

    def generate_image(self, grid: np.array, img_size: int = 512) -> Image:
        return generate_image_from_grid(grid, img_size)