import numpy as np
import heapq as hq
from collections import defaultdict

class Rectangle:
    def __init__(self, h, w, x, y, n, m):
        self.h, self.w = h, w
        self.x, self.y = x, y
        self.n, self.m = n, m
        self.mask = self.bit_mask()
        self.area = sum(self.mask)
    
    def bit_mask(self):
        mask = np.zeros((self.n, self.m), dtype=bool)
        mask[self.x : self.x + self.h, self.y : self.y + self.w] = True
        return mask.flatten()
    
    def __repr__(self):
        return f"Rectangle at ({self.x}, {self.y}) of size h = {self.h} and w = {self.w}\nMask : {self.mask}"

class Board:
    def __init__(self, rows, rectangles):
        self.n, self.m = len(rows), len(rows[0])
        mask = np.zeros((self.n, self.m), dtype=bool)
        self.frees = set()
        for i in range(self.n):
            for j in range(self.m):
                if rows[i][j] == " ":
                    mask[i][j] = True
        self.mask = mask.flatten()

        print(self.mask)
        self.frees = set([i for i, x in enumerate(self.mask) if not x])
        self.positions = self.get_rectangle_positions(rectangles)
    
    def show(self, mask):
        return "\n".join(
            "".join("." if b else "O" for b in mask[i * self.m : (i + 1) * self.m])
            for i in range(self.n)
        )
    
    def __repr__(self):
        pos_str = "\n".join(f"{k}: {v}" for k, v in self.positions.items())
        return f"Board\nMask :\n{self.show(self.mask)}\nFree positions :{self.frees}"
    
    def get_rectangle_positions(self, rects):
        pos = defaultdict(list)
        for free in self.frees:
            x, y = divmod(free, self.m)
            for idx, rect in enumerate(rects):
                h, w = rect 
                r = Rectangle(h, w, x, y, self.n, self.m)
                mask = r.mask
                area = r.area
                if not ((mask & self.mask).any()):
                    pos[free].append((idx, mask, area, 0))
                if h != w:
                    h, w = w, h
                    r = Rectangle(h, w, x, y, self.n, self.m)
                    mask = r.mask
                    area = r.area
                    if not ((mask & self.mask).any()):
                        pos[free].append((idx, mask, area, 1))
        return pos
    
    def solve(self):
        q = []
        hq.heapify(q)
        counter = 0  # Add a counter for tie-breaking
        
        for free in self.frees:
            for idx, mask, area, rot in self.positions[free]:
                if not ((mask & self.mask).any()):
                    new_mask = self.mask | mask  
                    # Add counter as tie-breaker to avoid comparing numpy arrays
                    ele = (-area, counter, new_mask, [(idx, free, rot)], set([idx]))
                    hq.heappush(q, ele)
                    counter += 1
        
        while len(q) > 0:
            _, _, mask, path, used = hq.heappop(q)  # Added extra _ for counter
            # print("used:", used)
            # print("path:", path)
            # print(self.show(mask))
            if mask.all():
                return path
            
            # Find first free position
            fst = None
            for i, val in enumerate(mask):
                if not val:
                    fst = i
                    break
            
            if fst is None:
                continue
                
            for idx, _mask, area, rot in self.positions[fst]:
                if idx not in used and not (_mask & mask).any():
                    new_path = path + [(idx, fst, rot)]
                    new_used = used | set([idx])
                    new_mask = mask | _mask
                    counter += 1
                    hq.heappush(
                        q,
                        (
                            -area,
                            counter,  # Add counter for tie-breaking
                            new_mask,
                            new_path,
                            new_used,
                        ),
                    )
        
        return None  # No solution found

board = [
    "            ",
    " 00000      ",
    " 00000      ",
    " 00000   00 ",
    "       000  ",
    "   00  000  ",
    " 0000 00    ",
    " 0000 00    ",
    " 00   000 0 ",
    "      000 0 ",
    "  0       0 ",
    "000         ",
]
pieces = [
    [1, 1],
    [1, 1],
    [1, 2],
    [1, 2],
    [1, 2],
    [1, 3],
    [1, 3],
    [1, 4],
    [1, 4],
    [2, 2],
    [2, 2],
    [2, 3],
    [2, 3],
    [2, 5],
]
# [(idx, free, rot)],
board = [
    "  000", 
    "00000"
        ]
'''
    "  200", 
    "31100"
    '''
pieces  = [[2,2], [2, 1], [1, 1], [1, 1]]
b = Board(board, pieces)
print(b)
r = b.solve()
print("r :", r)