from pathlib import Path
from PIL import Image
import numpy as np

def get_jigsaw_pieces(size):
    '''
    Load all pieces of the 4x4 jigsaw puzzle.

    size (int) :
        Should be 64 or 256, indicating side length of jigsaw puzzle
    '''

    # Location of pieces
    piece_dir = Path(__file__).parent / 'assets'

    # Helper function to load pieces as np arrays
    def load_pieces(path):
        '''
        Load a piece, from the given path, as a binary numpy array.
        Return a list of the "base" piece, and all four of its rotations.
        '''
        piece = Image.open(path)
        piece = np.array(piece)[:,:,0] // 255
        pieces = np.stack([np.rot90(piece, k=-i) for i in range(4)])
        return pieces

    # Load pieces and rotate to get 16 pieces, and cat
    pieces_corner = load_pieces(piece_dir / f'4x4/4x4_corner_{size}.png')
    pieces_inner = load_pieces(piece_dir / f'4x4/4x4_inner_{size}.png')
    pieces_edge1 = load_pieces(piece_dir / f'4x4/4x4_edge1_{size}.png')
    pieces_edge2 = load_pieces(piece_dir / f'4x4/4x4_edge2_{size}.png')
    pieces = np.concatenate([pieces_corner, pieces_inner, pieces_edge1, pieces_edge2])

    return pieces