from typing import List, Optional, Tuple, Union

import PIL.Image
import torch
from torchvision.prototype import features
from torchvision.prototype.features import BoundingBoxFormat, ColorSpace
from torchvision.transforms import functional_pil as _FP, functional_tensor as _FT


def get_dimensions_image_tensor(image: torch.Tensor) -> List[int]:
    chw = list(image.shape[-3:])
    ndims = len(chw)
    if ndims == 3:
        return chw
    elif ndims == 2:
        chw.insert(0, 1)
        return chw
    else:
        raise TypeError(f"Input tensor should have at least two dimensions, but got {ndims}")


get_dimensions_image_pil = _FP.get_dimensions


def get_dimensions(image: Union[features.ImageTypeJIT, features.VideoTypeJIT]) -> List[int]:
    if isinstance(image, torch.Tensor) and (
        torch.jit.is_scripting() or not isinstance(image, (features.Image, features.Video))
    ):
        return get_dimensions_image_tensor(image)
    elif isinstance(image, (features.Image, features.Video)):
        channels = image.num_channels
        height, width = image.spatial_size
        return [channels, height, width]
    else:
        return get_dimensions_image_pil(image)


def get_num_channels_image_tensor(image: torch.Tensor) -> int:
    chw = image.shape[-3:]
    ndims = len(chw)
    if ndims == 3:
        return chw[0]
    elif ndims == 2:
        return 1
    else:
        raise TypeError(f"Input tensor should have at least two dimensions, but got {ndims}")


get_num_channels_image_pil = _FP.get_image_num_channels


def get_num_channels_video(video: torch.Tensor) -> int:
    return get_num_channels_image_tensor(video)


def get_num_channels(image: Union[features.ImageTypeJIT, features.VideoTypeJIT]) -> int:
    if isinstance(image, torch.Tensor) and (
        torch.jit.is_scripting() or not isinstance(image, (features.Image, features.Video))
    ):
        return get_num_channels_image_tensor(image)
    elif isinstance(image, (features.Image, features.Video)):
        return image.num_channels
    else:
        return get_num_channels_image_pil(image)


# We changed the names to ensure it can be used not only for images but also videos. Thus, we just alias it without
# deprecating the old names.
get_image_num_channels = get_num_channels


def get_spatial_size_image_tensor(image: torch.Tensor) -> List[int]:
    hw = list(image.shape[-2:])
    ndims = len(hw)
    if ndims == 2:
        return hw
    else:
        raise TypeError(f"Input tensor should have at least two dimensions, but got {ndims}")


@torch.jit.unused
def get_spatial_size_image_pil(image: PIL.Image.Image) -> List[int]:
    width, height = _FP.get_image_size(image)
    return [height, width]


def get_spatial_size_video(video: torch.Tensor) -> List[int]:
    return get_spatial_size_image_tensor(video)


def get_spatial_size_mask(mask: torch.Tensor) -> List[int]:
    return get_spatial_size_image_tensor(mask)


@torch.jit.unused
def get_spatial_size_bounding_box(bounding_box: features.BoundingBox) -> List[int]:
    return list(bounding_box.spatial_size)


def get_spatial_size(inpt: features.InputTypeJIT) -> List[int]:
    if isinstance(inpt, torch.Tensor) and (torch.jit.is_scripting() or not isinstance(inpt, features._Feature)):
        return get_spatial_size_image_tensor(inpt)
    elif isinstance(inpt, (features.Image, features.Video, features.BoundingBox, features.Mask)):
        return list(inpt.spatial_size)
    else:
        return get_spatial_size_image_pil(inpt)  # type: ignore[no-any-return]


def get_num_frames_video(video: torch.Tensor) -> int:
    return video.shape[-4]


def get_num_frames(inpt: features.VideoTypeJIT) -> int:
    if isinstance(inpt, torch.Tensor) and (torch.jit.is_scripting() or not isinstance(inpt, features.Video)):
        return get_num_frames_video(inpt)
    elif isinstance(inpt, features.Video):
        return inpt.num_frames
    else:
        raise TypeError(f"The video should be a Tensor. Got {type(inpt)}")


def _xywh_to_xyxy(xywh: torch.Tensor) -> torch.Tensor:
    xyxy = xywh.clone()
    xyxy[..., 2:] += xyxy[..., :2]
    return xyxy


def _xyxy_to_xywh(xyxy: torch.Tensor) -> torch.Tensor:
    xywh = xyxy.clone()
    xywh[..., 2:] -= xywh[..., :2]
    return xywh


def _cxcywh_to_xyxy(cxcywh: torch.Tensor) -> torch.Tensor:
    cx, cy, w, h = torch.unbind(cxcywh, dim=-1)
    x1 = cx - 0.5 * w
    y1 = cy - 0.5 * h
    x2 = cx + 0.5 * w
    y2 = cy + 0.5 * h
    return torch.stack((x1, y1, x2, y2), dim=-1).to(cxcywh.dtype)


def _xyxy_to_cxcywh(xyxy: torch.Tensor) -> torch.Tensor:
    x1, y1, x2, y2 = torch.unbind(xyxy, dim=-1)
    cx = (x1 + x2) / 2
    cy = (y1 + y2) / 2
    w = x2 - x1
    h = y2 - y1
    return torch.stack((cx, cy, w, h), dim=-1).to(xyxy.dtype)


def convert_format_bounding_box(
    bounding_box: torch.Tensor, old_format: BoundingBoxFormat, new_format: BoundingBoxFormat
) -> torch.Tensor:
    if new_format == old_format:
        return bounding_box

    if old_format == BoundingBoxFormat.XYWH:
        bounding_box = _xywh_to_xyxy(bounding_box)
    elif old_format == BoundingBoxFormat.CXCYWH:
        bounding_box = _cxcywh_to_xyxy(bounding_box)

    if new_format == BoundingBoxFormat.XYWH:
        bounding_box = _xyxy_to_xywh(bounding_box)
    elif new_format == BoundingBoxFormat.CXCYWH:
        bounding_box = _xyxy_to_cxcywh(bounding_box)

    return bounding_box


def clamp_bounding_box(
    bounding_box: torch.Tensor, format: BoundingBoxFormat, spatial_size: Tuple[int, int]
) -> torch.Tensor:
    # TODO: Investigate if it makes sense from a performance perspective to have an implementation for every
    #  BoundingBoxFormat instead of converting back and forth
    xyxy_boxes = (
        bounding_box.clone()
        if format == BoundingBoxFormat.XYXY
        else convert_format_bounding_box(bounding_box, format, BoundingBoxFormat.XYXY)
    )
    xyxy_boxes[..., 0::2].clamp_(min=0, max=spatial_size[1])
    xyxy_boxes[..., 1::2].clamp_(min=0, max=spatial_size[0])
    return convert_format_bounding_box(xyxy_boxes, BoundingBoxFormat.XYXY, format)


def _split_alpha(image: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    return image[..., :-1, :, :], image[..., -1:, :, :]


def _strip_alpha(image: torch.Tensor) -> torch.Tensor:
    image, alpha = _split_alpha(image)
    if not torch.all(alpha == _FT._max_value(alpha.dtype)):
        raise RuntimeError(
            "Stripping the alpha channel if it contains values other than the max value is not supported."
        )
    return image


def _add_alpha(image: torch.Tensor, alpha: Optional[torch.Tensor] = None) -> torch.Tensor:
    if alpha is None:
        shape = list(image.shape)
        shape[-3] = 1
        alpha = torch.full(shape, _FT._max_value(image.dtype), dtype=image.dtype, device=image.device)
    return torch.cat((image, alpha), dim=-3)


def _gray_to_rgb(grayscale: torch.Tensor) -> torch.Tensor:
    repeats = [1] * grayscale.ndim
    repeats[-3] = 3
    return grayscale.repeat(repeats)


def _rgb_to_gray(image: torch.Tensor) -> torch.Tensor:
    r, g, b = image.unbind(dim=-3)
    l_img = (0.2989 * r).add_(g, alpha=0.587).add_(b, alpha=0.114)
    l_img = l_img.to(image.dtype).unsqueeze(dim=-3)
    return l_img


def convert_color_space_image_tensor(
    image: torch.Tensor, old_color_space: ColorSpace, new_color_space: ColorSpace
) -> torch.Tensor:
    if new_color_space == old_color_space:
        return image

    if old_color_space == ColorSpace.OTHER or new_color_space == ColorSpace.OTHER:
        raise RuntimeError(f"Conversion to or from {ColorSpace.OTHER} is not supported.")

    if old_color_space == ColorSpace.GRAY and new_color_space == ColorSpace.GRAY_ALPHA:
        return _add_alpha(image)
    elif old_color_space == ColorSpace.GRAY and new_color_space == ColorSpace.RGB:
        return _gray_to_rgb(image)
    elif old_color_space == ColorSpace.GRAY and new_color_space == ColorSpace.RGB_ALPHA:
        return _add_alpha(_gray_to_rgb(image))
    elif old_color_space == ColorSpace.GRAY_ALPHA and new_color_space == ColorSpace.GRAY:
        return _strip_alpha(image)
    elif old_color_space == ColorSpace.GRAY_ALPHA and new_color_space == ColorSpace.RGB:
        return _gray_to_rgb(_strip_alpha(image))
    elif old_color_space == ColorSpace.GRAY_ALPHA and new_color_space == ColorSpace.RGB_ALPHA:
        image, alpha = _split_alpha(image)
        return _add_alpha(_gray_to_rgb(image), alpha)
    elif old_color_space == ColorSpace.RGB and new_color_space == ColorSpace.GRAY:
        return _rgb_to_gray(image)
    elif old_color_space == ColorSpace.RGB and new_color_space == ColorSpace.GRAY_ALPHA:
        return _add_alpha(_rgb_to_gray(image))
    elif old_color_space == ColorSpace.RGB and new_color_space == ColorSpace.RGB_ALPHA:
        return _add_alpha(image)
    elif old_color_space == ColorSpace.RGB_ALPHA and new_color_space == ColorSpace.GRAY:
        return _rgb_to_gray(_strip_alpha(image))
    elif old_color_space == ColorSpace.RGB_ALPHA and new_color_space == ColorSpace.GRAY_ALPHA:
        image, alpha = _split_alpha(image)
        return _add_alpha(_rgb_to_gray(image), alpha)
    elif old_color_space == ColorSpace.RGB_ALPHA and new_color_space == ColorSpace.RGB:
        return _strip_alpha(image)
    else:
        raise RuntimeError(f"Conversion from {old_color_space} to {new_color_space} is not supported.")


_COLOR_SPACE_TO_PIL_MODE = {
    ColorSpace.GRAY: "L",
    ColorSpace.GRAY_ALPHA: "LA",
    ColorSpace.RGB: "RGB",
    ColorSpace.RGB_ALPHA: "RGBA",
}


@torch.jit.unused
def convert_color_space_image_pil(image: PIL.Image.Image, color_space: ColorSpace) -> PIL.Image.Image:
    old_mode = image.mode
    try:
        new_mode = _COLOR_SPACE_TO_PIL_MODE[color_space]
    except KeyError:
        raise ValueError(f"Conversion from {ColorSpace.from_pil_mode(old_mode)} to {color_space} is not supported.")

    if image.mode == new_mode:
        return image

    return image.convert(new_mode)


def convert_color_space_video(
    video: torch.Tensor, old_color_space: ColorSpace, new_color_space: ColorSpace
) -> torch.Tensor:
    return convert_color_space_image_tensor(video, old_color_space=old_color_space, new_color_space=new_color_space)


def convert_color_space(
    inpt: Union[features.ImageTypeJIT, features.VideoTypeJIT],
    color_space: ColorSpace,
    old_color_space: Optional[ColorSpace] = None,
) -> Union[features.ImageTypeJIT, features.VideoTypeJIT]:
    if isinstance(inpt, torch.Tensor) and (
        torch.jit.is_scripting() or not isinstance(inpt, (features.Image, features.Video))
    ):
        if old_color_space is None:
            raise RuntimeError(
                "In order to convert the color space of simple tensors, "
                "the `old_color_space=...` parameter needs to be passed."
            )
        return convert_color_space_image_tensor(inpt, old_color_space=old_color_space, new_color_space=color_space)
    elif isinstance(inpt, features.Image):
        output = convert_color_space_image_tensor(
            inpt.as_subclass(torch.Tensor), old_color_space=inpt.color_space, new_color_space=color_space
        )
        return features.Image.wrap_like(inpt, output, color_space=color_space)
    elif isinstance(inpt, features.Video):
        output = convert_color_space_video(
            inpt.as_subclass(torch.Tensor), old_color_space=inpt.color_space, new_color_space=color_space
        )
        return features.Video.wrap_like(inpt, output, color_space=color_space)
    else:
        return convert_color_space_image_pil(inpt, color_space)


def _num_value_bits(dtype: torch.dtype) -> int:
    if dtype == torch.uint8:
        return 8
    elif dtype == torch.int8:
        return 7
    elif dtype == torch.int16:
        return 15
    elif dtype == torch.int32:
        return 31
    elif dtype == torch.int64:
        return 63
    else:
        raise TypeError(f"Number of value bits is only defined for integer dtypes, but got {dtype}.")


def convert_dtype_image_tensor(image: torch.Tensor, dtype: torch.dtype = torch.float) -> torch.Tensor:
    if image.dtype == dtype:
        return image

    float_input = image.is_floating_point()
    if torch.jit.is_scripting():
        # TODO: remove this branch as soon as `dtype.is_floating_point` is supported by JIT
        float_output = torch.tensor(0, dtype=dtype).is_floating_point()
    else:
        float_output = dtype.is_floating_point

    if float_input:
        # float to float
        if float_output:
            return image.to(dtype)

        # float to int
        if (image.dtype == torch.float32 and dtype in (torch.int32, torch.int64)) or (
            image.dtype == torch.float64 and dtype == torch.int64
        ):
            raise RuntimeError(f"The conversion from {image.dtype} to {dtype} cannot be performed safely.")

        # For data in the range `[0.0, 1.0]`, just multiplying by the maximum value of the integer range and converting
        # to the integer dtype  is not sufficient. For example, `torch.rand(...).mul(255).to(torch.uint8)` will only
        # be `255` if the input is exactly `1.0`. See https://github.com/pytorch/vision/pull/2078#issuecomment-612045321
        # for a detailed analysis.
        # To mitigate this, we could round before we convert to the integer dtype, but this is an extra operation.
        # Instead, we can also multiply by the maximum value plus something close to `1`. See
        # https://github.com/pytorch/vision/pull/2078#issuecomment-613524965 for details.
        eps = 1e-3
        max_value = float(_FT._max_value(dtype))
        # We need to scale first since the conversion would otherwise turn the input range `[0.0, 1.0]` into the
        # discrete set `{0, 1}`.
        return image.mul(max_value + 1.0 - eps).to(dtype)
    else:
        # int to float
        if float_output:
            return image.to(dtype).div_(_FT._max_value(image.dtype))

        # int to int
        num_value_bits_input = _num_value_bits(image.dtype)
        num_value_bits_output = _num_value_bits(dtype)

        if num_value_bits_input > num_value_bits_output:
            return image.bitwise_right_shift(num_value_bits_input - num_value_bits_output).to(dtype)
        else:
            # The bitshift kernel is not vectorized
            #  https://github.com/pytorch/pytorch/blob/703c19008df4700b6a522b0ae5c4b6d5ffc0906f/aten/src/ATen/native/cpu/BinaryOpsKernel.cpp#L315-L322
            #  This results in the multiplication actually being faster.
            # TODO: If the bitshift kernel is optimized in core, replace the computation below with
            #  `image.to(dtype).bitwise_left_shift_(num_value_bits_output - num_value_bits_input)`
            max_value_input = float(_FT._max_value(dtype))
            max_value_output = float(_FT._max_value(image.dtype))
            factor = int((max_value_input + 1) // (max_value_output + 1))
            return image.to(dtype).mul_(factor)


# We changed the name to align it with the new naming scheme. Still, `convert_image_dtype` is
# prevalent and well understood. Thus, we just alias it without deprecating the old name.
convert_image_dtype = convert_dtype_image_tensor


def convert_dtype_video(video: torch.Tensor, dtype: torch.dtype = torch.float) -> torch.Tensor:
    return convert_dtype_image_tensor(video, dtype)


def convert_dtype(
    inpt: Union[features.ImageTypeJIT, features.VideoTypeJIT], dtype: torch.dtype = torch.float
) -> torch.Tensor:
    if isinstance(inpt, torch.Tensor) and (
        torch.jit.is_scripting() or not isinstance(inpt, (features.Image, features.Video))
    ):
        return convert_dtype_image_tensor(inpt, dtype)
    elif isinstance(inpt, features.Image):
        output = convert_dtype_image_tensor(inpt.as_subclass(torch.Tensor), dtype)
        return features.Image.wrap_like(inpt, output)
    else:  # isinstance(inpt, features.Video):
        output = convert_dtype_video(inpt.as_subclass(torch.Tensor), dtype)
        return features.Video.wrap_like(inpt, output)
