from PIL import Image
from pyzbar.pyzbar import decode
import itertools
from math import perm
from tqdm import tqdm
import numpy as np
import sys

NUM_PIECES_WIDTH = 4
NUM_PIECES_HEIGHT = 3
BOX_SIZE = 2
TILE = 78
# 7x7 finder, upscaled to 14x14 (black=1, white=0)
TEMPLATE = np.kron(np.array([
    [1,1,1,1,1,1,1],
    [1,0,0,0,0,0,1],
    [1,0,1,1,1,0,1],
    [1,0,1,1,1,0,1],
    [1,0,1,1,1,0,1],
    [1,0,0,0,0,0,1],
    [1,1,1,1,1,1,1],
], dtype=np.uint8), np.ones((BOX_SIZE, BOX_SIZE), np.uint8))
H, W = TEMPLATE.shape  # 14x14


def decodeQRCode(image : Image.Image) -> str | None:
    # Decode the QR code
    decoded_objects = decode(image)

    # Check if any QR code was found
    if decoded_objects:
        # Iterate through all the detected QR codes and print their data
        for obj in decoded_objects:
            return obj.data.decode("utf-8")
    else:
        return None


def cut_pieces(image):
    # Get the size of the image
    width, height = image.size

    # Calculate the size of each grid cell
    cell_width = width // NUM_PIECES_WIDTH 
    cell_height = height // NUM_PIECES_HEIGHT 

    pieces = []

    # Split the original image into 12 pieces
    for i in range(NUM_PIECES_WIDTH * NUM_PIECES_HEIGHT):
        # Calculate the coordinates of the current grid cell
        x1 = (i % NUM_PIECES_WIDTH) * cell_width
        y1 = (i // NUM_PIECES_WIDTH) * cell_height
        x2 = x1 + cell_width
        y2 = y1 + cell_height

        # Crop the corresponding part of the original image
        cropped_piece = image.crop((x1, y1, x2, y2))
        pieces.append(cropped_piece)

    return pieces


def to_bw01(img):
    # Force 1-bit, then map black->1, white->0
    a = np.array(img.convert("1"))  # {0,255}
    return (a == 0).astype(np.uint8)


def finder_corner_exact(tile_bw):
    crops = {
        'TL': tile_bw[0:H,            0:W            ],
        'TR': tile_bw[0:H,            TILE-W:TILE    ],
        'BL': tile_bw[TILE-H:TILE,    0:W            ],
        'BR': tile_bw[TILE-H:TILE,    TILE-W:TILE    ],
    }
    for corner, crop in crops.items():
        if np.array_equal(crop, TEMPLATE):
            return corner
    return None  # no exact finder in this tile


def three_tiles_with_finders(tiles):
    hits = []
    for tile in tiles:
        corner = finder_corner_exact(to_bw01(tile))
        if corner is not None:
            hits.append((tile, corner))
            if len(hits) == 3:
                break
    return hits


if __name__ == "__main__":
    img = Image.open("./scrambled_test_qr_code.png")
    new_image = Image.new("RGB", (img.height, img.height))

    pieces = cut_pieces(img)
    corners = three_tiles_with_finders(pieces)

    for tile, corner in corners:
        if corner == 'TL':
            new_image.paste(tile, (0,0))
        elif corner == 'TR':
            new_image.paste(tile, (156,0))
        elif corner == 'BL':
            new_image.paste(tile, (0,156))
        else:
            sys.exit()


    pieces = [piece for piece in pieces if piece not in [tile for tile, _ in corners]]
    assert len(pieces) == 9
    new_image.show()
    possible_combos = itertools.permutations(pieces, 6)
    num_possible_combos = perm(len(pieces), 6)
    print(f"Testing {num_possible_combos} combinations")

    for a,b,c,d,e,f in tqdm(possible_combos, total=num_possible_combos):
        new_image.paste(a, (78, 0))
        new_image.paste(b, (0, 78))
        new_image.paste(c, (78, 78))
        new_image.paste(d, (156, 78))
        new_image.paste(e, (78, 156))
        new_image.paste(f, (156, 156))
        output = decodeQRCode(new_image)
        if output:
            print(output)

    print(decodeQRCode(img))