import numpy as np

three_in_a_row = np.array([
    [1, 1, 1]
])

l_shape = np.array([
    [1, 0],
    [1, 0],
    [1, 1]
])

t_shape = np.array([
    [1, 1, 1],
    [0, 1, 0],
])

square = np.array([
    [1, 1],
    [1, 1]
])

smaller_l_shape = np.array([
    [1, 0],
    [1, 1]
])

tetris_z = np.array([
    [1, 1, 0],
    [0, 1, 1]
])

all_pieces = {
    "Consecutive 3s": three_in_a_row,
    "L": l_shape,
    "T": t_shape,
    "Square": square,
    "Smaller L": smaller_l_shape,
    "Z": tetris_z
}

def sanity_check():
    """Check all possible ways to position the pieces(including 3d rotation)"""
    for piece_name, piece in all_pieces.items():
        all_placements = []

        # Check original rotations
        for _ in range(4):
            piece = np.rot90(piece)
            if piece.tolist() not in all_placements:
                all_placements.append(piece.tolist())

        # Check flipped rotations
        flipped_piece = np.flip(piece, axis=0)
        for _ in range(4):
            flipped_piece = np.rot90(flipped_piece)
            if flipped_piece.tolist() not in all_placements:
                all_placements.append(flipped_piece.tolist())
        
        print(f"Piece: {piece_name}")
        print(f"Distinct Placements (including flips and rotations): {len(all_placements)}")
        print("--------------------")

def count_combinations_on_matrix(matrix_size):
    """Count all possible combinations of pieces on a matrix of given size."""
    matrix = np.zeros(matrix_size, dtype=int)
    total_combinations = 0

    for piece_name, piece in all_pieces.items():
        piece_rows, piece_cols = piece.shape
        placements = []

        # Generate all distinct placements of the piece
        for _ in range(4):
            piece = np.rot90(piece)
            if piece.tolist() not in placements:
                placements.append(piece.tolist())

        flipped_piece = np.flip(piece, axis=0)
        for _ in range(4):
            flipped_piece = np.rot90(flipped_piece)
            if flipped_piece.tolist() not in placements:
                placements.append(flipped_piece.tolist())

        # Count valid placements on the matrix
        for placement in placements:
            rows, cols = len(placement), len(placement[0])
            for i in range(matrix_size[0] - rows + 1):
                for j in range(matrix_size[1] - cols + 1):
                    sub_matrix = matrix[i:i + rows, j:j + cols]
                    if not sub_matrix.any():  # Check if the space is empty
                        total_combinations += 1

    print(f"Total combinations on {matrix_size[0]}x{matrix_size[1]} matrix: {total_combinations}")


if __name__ == "__main__":
    sanity_check()
    count_combinations_on_matrix((5, 5))