import torch

class PatchShuffleTransform(object):
    def __init__(self, patch_size=32):
        self.patch_size = patch_size

    def __call__(self, tensor):
        c, h, w = tensor.shape
        if h % self.patch_size != 0 or w % self.patch_size != 0:
            raise ValueError("Image dimensions must be divisible by patch_size")
        # Divide the image into patches
        tensor = tensor.unfold(1, self.patch_size, self.patch_size).unfold(2, self.patch_size, self.patch_size)
        num_patches_h = tensor.size(1)
        num_patches_w = tensor.size(2)
        patches = tensor.contiguous().view(c, -1, self.patch_size, self.patch_size)
        # Shuffle patches
        idx = torch.randperm(patches.size(1))
        patches = patches[:, idx, :, :]
        # Reshape back to image
        patches = patches.view(c, num_patches_h, num_patches_w, self.patch_size, self.patch_size)
        patches = patches.permute(0, 1, 3, 2, 4).contiguous().view(c, h, w)
        return patches

    def __repr__(self):
        return f"{self.__class__.__name__}(patch_size={self.patch_size})"
