import random
from itertools import product
from typing import Tuple, Set, Union, Any, List, NamedTuple
from pprint import pprint


class Cell(NamedTuple):
    x: int
    y: int


class Wall(NamedTuple):
    first_cell: Cell
    second_cell: Cell


class Maze:
    def __init__(self, width: int, height: int):
        self.distance_map = dict()
        self.cells = set(Cell(*c) for c in product(range(width), range(height)))
        self.maze = dict()
        for cell in self.cells:
            self.maze[cell] = self.get_neighbours(cell)
        self.walls = self.get_walls()
        self.generate_maze()
        self.remove_blocked_links()
        self.build_map()
        self.visited = set()
        self.dead_ends = dict()
        self.dead_end_limit = 5

    def check_dead_end(self, cell: Cell) -> bool:
        # return true is new dead end
        if cell in self.dead_ends:
            return False
        if len(self.maze[cell]) == 1:
            self.dead_ends[cell] = 0
            return True


    def build_map(self):
        distance = 0
        first_cell = max(self.cells)
        visited = set()
        visited.add(first_cell)
        unvisited = self.cells.copy()
        unvisited.remove(first_cell)
        self.distance_map[first_cell] = 0
        current_group = self.maze[first_cell]

        while unvisited:
            next_group = set()
            distance += 1
            for node in current_group:
                self.distance_map[node] = distance
                next_group.update(self.maze[node])
                visited.add(node)
                unvisited.remove(node)
            current_group = next_group.difference(visited)

    def get_neighbours(self, cell: Cell) -> List[Cell]:
        diffs = ((0, 1), (0, -1), (1, 0), (-1, 0))
        return list(Cell(cell.x + x, cell.y + y) for x, y in diffs if Cell(cell.x + x, cell.y + y) in self.cells)

    def get_walls(self) -> Set[Wall]:
        walls = set()
        for cell in self.cells:
            walls.update(set(Wall(*sorted([cell, c])) for c in self.maze[cell]))
        return walls

    def remove_blocked_links(self):
        for wall in self.walls:
            cell1, cell2 = wall
            self.maze[cell1].remove(cell2)
            self.maze[cell2].remove(cell1)

    def generate_maze(self) -> None:
        visited = set()
        unvisited = list(self.cells)
        node = random.choice(unvisited)
        visited.add(node)
        unvisited.remove(node)
        while unvisited:
            random_neighbour = random.choice(list(self.maze[node]))
            if random_neighbour not in visited:
                self.walls.remove(Wall(*sorted([node, random_neighbour])))
                visited.add(random_neighbour)
                try:
                    unvisited.remove(random_neighbour)
                except ValueError:
                    pass
            node = random_neighbour


if __name__ == "__main__":
    maze = Maze(3, 3)
    pprint(maze.walls)