from typing import Tuple
import torch


def _sobel_kernel(spatial_size: int) -> Tuple[torch.Tensor, ...]:
    """
        Constructs the sobel kernel for 2D or 3D kernel sizes.

        Args:
            spatial_size (int): The spatial size of the sobel kernel.

        Raises:
            ValueError when spatial_size is not 2 or 3.
    """
    if spatial_size == 2:
        g_x = (
            torch.Tensor([[+1, 0, -1], [+2, 0, -2], [+1, 0, -1]])
            .unsqueeze(0)
            .unsqueeze(0)
        )

        g_y = (
            torch.Tensor([[+1, +2, +1], [0, 0, 0], [-1, -2, 1]])
            .unsqueeze(0)
            .unsqueeze(0)
        )

        return g_x, g_y
    elif spatial_size == 3:
        g_x = (
            torch.Tensor(
                [
                    [[1, 0, -1], [1, 0, -1], [1, 0, -1]],
                    [[1, 0, -1], [2, 0, -2], [1, 0, -1]],
                    [[1, 0, -1], [1, 0, -1], [1, 0, -1]],
                ]
            )
            .unsqueeze(0)
            .unsqueeze(0)
        )

        g_y = (
            torch.Tensor(
                [
                    [[1, 1, 1], [0, 0, 0], [-1, -1, -1]],
                    [[1, 2, 1], [0, 0, 0], [-1, -2, -1]],
                    [[1, 1, 1], [0, 0, 0], [-1, -1, -1]],
                ]
            )
            .unsqueeze(0)
            .unsqueeze(0)
        )

        g_z = (
            torch.Tensor(
                [
                    [[1, 1, 1], [1, 2, 1], [1, 1, 1]],
                    [[0, 0, 0], [0, 0, 0], [0, 0, 0]],
                    [[-1, -1, -1], [-1, -2, -1], [-1, -1, -1]],
                ]
            )
            .unsqueeze(0)
            .unsqueeze(0)
        )

        return g_x, g_y, g_z
    else:
        raise ValueError(
            f"Sobel operator is defined only for spatial sizes 2 and 3, but given {spatial_size}."
        )