"""PyTorch tensor utilities for image processing."""

import numpy as np
import torch
from einops import rearrange
from PIL import Image
from torchvision.transforms.functional import to_pil_image


def numpy_to_torch(img: np.ndarray) -> torch.Tensor:
    """
    Convert numpy image to torch tensor.

    For 3D arrays (H, W, C), permutes to (C, H, W).
    For 2D arrays (H, W), passes through unchanged.

    Args:
        img: Input numpy array of shape (H, W, C) or (H, W)

    Returns:
        Torch tensor of shape (C, H, W) or (H, W)
    """
    t = torch.from_numpy(img)
    if t.ndim == 3:
        t = t.permute(2, 0, 1)
    return t


def normalize_uint8_to_neg1_1(x: torch.Tensor) -> torch.Tensor:
    """
    Normalize uint8 image tensor from [0, 255] to [-1, 1] range.

    Args:
        x: Input tensor with values in [0, 255]

    Returns:
        Normalized tensor with values in [-1, 1]
    """
    return x / 127.5 - 1.0


def _neg1_1_to_0_1(normed_img: torch.Tensor) -> torch.Tensor:
    """Convert [-1, 1] normalized tensor to [0, 1] range."""
    return (normed_img + 1) * 0.5


def tensor_to_pil(img: torch.Tensor, unnormalize: bool = False) -> Image.Image:
    """
    Convert PyTorch tensor to PIL Image.

    Args:
        img: Input tensor of shape (C, H, W)
        unnormalize: If True, convert from [-1, 1] to [0, 1] range first

    Returns:
        PIL Image
    """
    if unnormalize:
        img = _neg1_1_to_0_1(img)
    return to_pil_image(img)


def unpack_images(x: torch.Tensor, patch_size: int = 2) -> torch.Tensor:
    """
    Unpack image patches back to full images.

    Used after transformer processing to convert patch representations
    back to spatial images.

    Args:
        x: Tensor of shape (batch_size, channels * patch_size^2, h, w)
        patch_size: Size of patches used during packing

    Returns:
        Tensor of shape (batch_size, channels, h * patch_size, w * patch_size)
    """
    return rearrange(x, "b (c p1 p2) h w -> b c (h p1) (w p2)", p1=patch_size, p2=patch_size)