import torch
import torch.nn.functional as F
import math

class CalculateTilingSpec_Tiled:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "width": ("INT", {"default": 512, "min": 1, "max": 16384}),
                "height": ("INT", {"default": 512, "min": 1, "max": 16384}),
                "rows": ("INT", {"default": 2, "min": 1, "max": 100}),
                "columns": ("INT", {"default": 2, "min": 1, "max": 100}),
                "overlap": ("INT", {"default": 64, "min": 0, "max": 1024}),
            }
        }

    RETURN_TYPES = ("*", "INT", "INT")
    RETURN_NAMES = ("tiling_spec", "tile_width", "tile_height")
    FUNCTION = "calculate"
    CATEGORY = "SplitImage"

    def calculate(self, width, height, rows, columns, overlap):
        tile_w = math.ceil((width + (columns - 1) * overlap) / columns)
        tile_h = math.ceil((height + (rows - 1) * overlap) / rows)
        
        pad_w = (columns * tile_w - (columns - 1) * overlap) - width
        pad_h = (rows * tile_h - (rows - 1) * overlap) - height
        
        tiling_spec = {
            "rows": rows,
            "columns": columns,
            "overlap": overlap,
            "tile_w": tile_w,
            "tile_h": tile_h,
            "orig_w": width,
            "orig_h": height,
            "padded_w": width + pad_w,
            "padded_h": height + pad_h
        }
        
        return (tiling_spec, tile_w, tile_h)

class SplitImage_Tiled:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),
                "rows": ("INT", {"default": 2, "min": 1, "max": 100}),
                "columns": ("INT", {"default": 2, "min": 1, "max": 100}),
                "overlap": ("INT", {"default": 64, "min": 0, "max": 1024}),
            }
        }

    RETURN_TYPES = ("IMAGE", "INT", "INT", "INT", "*")
    RETURN_NAMES = ("tiles", "rows", "columns", "overlap", "tiling_spec")
    OUTPUT_IS_LIST = (True, False, False, False, False)
    FUNCTION = "split"
    CATEGORY = "SplitImage"

    def split(self, image, rows, columns, overlap):
        B, H, W, C = image.shape
        
        tile_w = math.ceil((W + (columns - 1) * overlap) / columns)
        tile_h = math.ceil((H + (rows - 1) * overlap) / rows)
        
        pad_w = (columns * tile_w - (columns - 1) * overlap) - W
        pad_h = (rows * tile_h - (rows - 1) * overlap) - H
        
        if pad_w > 0 or pad_h > 0:
            image_padded = image.permute(0, 3, 1, 2)
            image_padded = F.pad(image_padded, (0, pad_w, 0, pad_h), mode='constant', value=0)
            image = image_padded.permute(0, 2, 3, 1)
            _, H, W, _ = image.shape

        tiles = []
        for r in range(rows):
            for c in range(columns):
                x1 = c * (tile_w - overlap)
                y1 = r * (tile_h - overlap)
                x2 = x1 + tile_w
                y2 = y1 + tile_h
                
                tile = image[:, y1:y2, x1:x2, :]
                # Store each tile as a single-image tensor in the list
                for i in range(tile.shape[0]):
                    tiles.append(tile[i:i+1])
        
        tiling_spec = {
            "rows": rows,
            "columns": columns,
            "overlap": overlap,
            "tile_w": tile_w,
            "tile_h": tile_h,
            "orig_w": W - pad_w,
            "orig_h": H - pad_h,
            "padded_w": W,
            "padded_h": H
        }
        
        return (tiles, rows, columns, overlap, tiling_spec)