import torch


class GridSlice:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "columns": ("INT", {"default": 8, "min": 1, "max": 64}),
                "rows": ("INT", {"default": 4, "min": 1, "max": 64}),
                "ignore_last_column": ("BOOLEAN", {"default": False}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "slice_image"
    CATEGORY = "Game Development/Image Adjustments"

    def slice_image(self, image, columns, rows, ignore_last_column):
        _, h, w, _ = image.shape

        frame_w = w // columns
        frame_h = h // rows
        columns_to_slice = (
            columns - 1 if ignore_last_column and columns > 1 else columns
        )

        frames = []
        for r in range(rows):
            for c in range(columns_to_slice):
                y1 = r * frame_h
                y2 = (r + 1) * frame_h
                x1 = c * frame_w
                x2 = (c + 1) * frame_w

                frame = image[:, y1:y2, x1:x2, :]
                frames.append(frame)

        batch_frames = torch.cat(frames, dim=0)
        return (batch_frames,)