"""只用 PyTorch reshape/permute 完成 ViT patchify 与逆变换。"""
import torch


def patchify(images: torch.Tensor, patch: int) -> torch.Tensor:
    # images: (B,C,H,W) -> (B,N,C*P*P)
    b, c, h, w = images.shape
    if h % patch or w % patch:
        raise ValueError("height and width must be divisible by patch size")
    x = images.reshape(b, c, h // patch, patch, w // patch, patch)
    return x.permute(0, 2, 4, 1, 3, 5).reshape(b, -1, c * patch * patch)


def unpatchify(tokens: torch.Tensor, channels: int, height: int, width: int, patch: int):
    b = tokens.shape[0]
    x = tokens.reshape(b, height // patch, width // patch, channels, patch, patch)
    return x.permute(0, 3, 1, 4, 2, 5).reshape(b, channels, height, width)


if __name__ == "__main__":
    image = torch.arange(2 * 3 * 8 * 8).reshape(2, 3, 8, 8)
    tokens = patchify(image, 4)
    restored = unpatchify(tokens, 3, 8, 8, 4)
    assert tokens.shape == (2, 4, 48)
    assert torch.equal(image, restored)