from typing import Union

import torch
import torch.nn as nn

from kornia.color.hsv import rgb_to_hsv, hsv_to_rgb
from kornia.constants import pi


def adjust_saturation_raw(input: torch.Tensor, saturation_factor: Union[float, torch.Tensor]) -> torch.Tensor:
    r"""Adjust color saturation of an image. Expecting input to be in hsv format already.

    See :class:`~kornia.color.AdjustSaturation` for details.
    """

    if not torch.is_tensor(input):
        raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}")

    if not isinstance(saturation_factor, (float, torch.Tensor,)):
        raise TypeError(f"The saturation_factor should be a float number or torch.Tensor."
                        f"Got {type(saturation_factor)}")

    if isinstance(saturation_factor, float):
        saturation_factor = torch.tensor([saturation_factor])

    saturation_factor = saturation_factor.to(input.device).to(input.dtype)

    if (saturation_factor < 0).any():
        raise ValueError(f"Saturation factor must be non-negative. Got {saturation_factor}")

    for _ in input.shape[1:]:
        saturation_factor = torch.unsqueeze(saturation_factor, dim=-1)

    # unpack the hsv values
    h, s, v = torch.chunk(input, chunks=3, dim=-3)

    # transform the hue value and appl module
    s_out: torch.Tensor = torch.clamp(s * saturation_factor, min=0, max=1)

    # pack back back the corrected hue
    out: torch.Tensor = torch.cat([h, s_out, v], dim=-3)

    return out


def adjust_saturation(input: torch.Tensor, saturation_factor: Union[float, torch.Tensor]) -> torch.Tensor:
    r"""Adjust color saturation of an image.

    See :class:`~kornia.color.AdjustSaturation` for details.
    """

    # convert the rgb image to hsv
    x_hsv: torch.Tensor = rgb_to_hsv(input)

    # perform the conversion
    x_adjusted: torch.Tensor = adjust_saturation_raw(x_hsv, saturation_factor)

    # convert back to rgb
    out: torch.Tensor = hsv_to_rgb(x_adjusted)

    return out


def adjust_hue_raw(input: torch.Tensor, hue_factor: Union[float, torch.Tensor]) -> torch.Tensor:
    r"""Adjust hue of an image. Expecting input to be in hsv format already.

    See :class:`~kornia.color.AdjustHue` for details.
    """

    if not torch.is_tensor(input):
        raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}")

    if not isinstance(hue_factor, (float, torch.Tensor)):
        raise TypeError(f"The hue_factor should be a float number or torch.Tensor in the range between"
                        f" [-PI, PI]. Got {type(hue_factor)}")

    if isinstance(hue_factor, float):
        hue_factor = torch.tensor([hue_factor])

    hue_factor = hue_factor.to(input.device).to(input.dtype)

    if ((hue_factor < -pi) | (hue_factor > pi)).any():
        raise ValueError(f"Hue-factor must be in the range [-PI, PI]. Got {hue_factor}")

    for _ in input.shape[1:]:
        hue_factor = torch.unsqueeze(hue_factor, dim=-1)

    # unpack the hsv values
    h, s, v = torch.chunk(input, chunks=3, dim=-3)

    # transform the hue value and appl module
    divisor: float = 2 * pi.item()
    h_out: torch.Tensor = torch.fmod(h + hue_factor, divisor)

    # pack back back the corrected hue
    out: torch.Tensor = torch.cat([h_out, s, v], dim=-3)

    return out


def adjust_hue(input: torch.Tensor, hue_factor: Union[float, torch.Tensor]) -> torch.Tensor:
    r"""Adjust hue of an image.

    See :class:`~kornia.color.AdjustHue` for details.
    """

    # convert the rgb image to hsv
    x_hsv: torch.Tensor = rgb_to_hsv(input)

    # perform the conversion
    x_adjusted: torch.Tensor = adjust_hue_raw(x_hsv, hue_factor)

    # convert back to rgb
    out: torch.Tensor = hsv_to_rgb(x_adjusted)

    return out


def adjust_gamma(input: torch.Tensor, gamma: Union[float, torch.Tensor],
                 gain: Union[float, torch.Tensor] = 1.) -> torch.Tensor:
    r"""Perform gamma correction on an image.

    See :class:`~kornia.color.AdjustGamma` for details.
    """

    if not torch.is_tensor(input):
        raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}")

    if not isinstance(gamma, (float, torch.Tensor)):
        raise TypeError(f"The gamma should be a positive float or torch.Tensor. Got {type(gamma)}")

    if not isinstance(gain, (float, torch.Tensor)):
        raise TypeError(f"The gain should be a positive float or torch.Tensor. Got {type(gain)}")

    if isinstance(gamma, float):
        gamma = torch.tensor([gamma])

    if isinstance(gain, float):
        gain = torch.tensor([gain])

    gamma = gamma.to(input.device).to(input.dtype)
    gain = gain.to(input.device).to(input.dtype)

    if (gamma < 0.0).any():
        raise ValueError(f"Gamma must be non-negative. Got {gamma}")

    if (gain < 0.0).any():
        raise ValueError(f"Gain must be non-negative. Got {gain}")

    for _ in input.shape[1:]:
        gamma = torch.unsqueeze(gamma, dim=-1)
        gain = torch.unsqueeze(gain, dim=-1)

    # Apply the gamma correction
    x_adjust: torch.Tensor = gain * torch.pow(input, gamma)

    # Truncate between pixel values
    out: torch.Tensor = torch.clamp(x_adjust, 0.0, 1.0)

    return out


def adjust_contrast(input: torch.Tensor,
                    contrast_factor: Union[float, torch.Tensor]) -> torch.Tensor:
    r"""Adjust Contrast of an image.

    See :class:`~kornia.color.AdjustContrast` for details.
    """

    if not torch.is_tensor(input):
        raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}")

    if not isinstance(contrast_factor, (float, torch.Tensor,)):
        raise TypeError(f"The factor should be either a float or torch.Tensor. "
                        f"Got {type(contrast_factor)}")

    if isinstance(contrast_factor, float):
        contrast_factor = torch.tensor([contrast_factor])

    contrast_factor = contrast_factor.to(input.device).to(input.dtype)

    if (contrast_factor < 0).any():
        raise ValueError(f"Contrast factor must be non-negative. Got {contrast_factor}")

    for _ in input.shape[1:]:
        contrast_factor = torch.unsqueeze(contrast_factor, dim=-1)

    # Apply contrast factor to each channel
    x_adjust: torch.Tensor = input * contrast_factor

    # Truncate between pixel values
    out: torch.Tensor = torch.clamp(x_adjust, 0.0, 1.0)

    return out


def adjust_brightness(input: torch.Tensor,
                      brightness_factor: Union[float, torch.Tensor]) -> torch.Tensor:
    r"""Adjust Brightness of an image.

    See :class:`~kornia.color.AdjustBrightness` for details.
    """

    if not torch.is_tensor(input):
        raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}")

    if not isinstance(brightness_factor, (float, torch.Tensor,)):
        raise TypeError(f"The factor should be either a float or torch.Tensor. "
                        f"Got {type(brightness_factor)}")

    if isinstance(brightness_factor, float):
        brightness_factor = torch.tensor([brightness_factor])

    brightness_factor = brightness_factor.to(input.device).to(input.dtype)

    for _ in input.shape[1:]:
        brightness_factor = torch.unsqueeze(brightness_factor, dim=-1)

    # Apply brightness factor to each channel
    x_adjust: torch.Tensor = input + brightness_factor

    # Truncate between pixel values
    out: torch.Tensor = torch.clamp(x_adjust, 0.0, 1.0)

    return out


class AdjustSaturation(nn.Module):
    r"""Adjust color saturation of an image.

    The input image is expected to be an RGB image in the range of [0, 1].

    Args:
        input (torch.Tensor): Image/Tensor to be adjusted in the shape of (\*, N).
        saturation_factor (float):  How much to adjust the saturation. 0 will give a black
        and white image, 1 will give the original image while 2 will enhance the saturation
        by a factor of 2.

    Returns:
        torch.Tensor: Adjusted image.
    """

    def __init__(self, saturation_factor: Union[float, torch.Tensor]) -> None:
        super(AdjustSaturation, self).__init__()
        self.saturation_factor: Union[float, torch.Tensor] = saturation_factor

    def forward(self, input: torch.Tensor) -> torch.Tensor:  # type: ignore
        return adjust_saturation(input, self.saturation_factor)


class AdjustHue(nn.Module):
    r"""Adjust hue of an image.

    The input image is expected to be an RGB image in the range of [0, 1].

    Args:
        input (torch.Tensor): Image/Tensor to be adjusted in the shape of (\*, N).
        hue_factor (float): How much to shift the hue channel. Should be in [-PI, PI]. PI
          and -PI give complete reversal of hue channel in HSV space in positive and negative
          direction respectively. 0 means no shift. Therefore, both -PI and PI will give an
          image with complementary colors while 0 gives the original image.

    Returns:
        torch.Tensor: Adjusted image.
    """

    def __init__(self, hue_factor: Union[float, torch.Tensor]) -> None:
        super(AdjustHue, self).__init__()
        self.hue_factor: Union[float, torch.Tensor] = hue_factor

    def forward(self, input: torch.Tensor) -> torch.Tensor:  # type: ignore
        return adjust_hue(input, self.hue_factor)


class AdjustGamma(nn.Module):
    r"""Perform gamma correction on an image.

    The input image is expected to be in the range of [0, 1].

    Args:
        input (torch.Tensor): Image/Tensor to be adjusted in the shape of (\*, N).
        gamma (float): Non negative real number, same as γ\gammaγ in the equation.
          gamma larger than 1 make the shadows darker, while gamma smaller than 1 make
          dark regions lighter.
        gain (float, optional): The constant multiplier. Default 1.

    Returns:
        torch.Tensor: Adjusted image.
    """

    def __init__(self, gamma: Union[float, torch.Tensor], gain: Union[float, torch.Tensor] = 1.) -> None:
        super(AdjustGamma, self).__init__()
        self.gamma: Union[float, torch.Tensor] = gamma
        self.gain: Union[float, torch.Tensor] = gain

    def forward(self, input: torch.Tensor) -> torch.Tensor:  # type: ignore
        return adjust_gamma(input, self.gamma, self.gain)


class AdjustContrast(nn.Module):
    r"""Adjust Contrast of an image. This implementation aligns OpenCV, not PIL. Hence,
    the output differs from TorchVision.

    The input image is expected to be in the range of [0, 1].

    Args:
        input (torch.Tensor): Image to be adjusted in the shape of (\*, N).
        contrast_factor (Union[float, torch.Tensor]): Contrast adjust factor per element
          in the batch. 0 generates a compleatly black image, 1 does not modify
          the input image while any other non-negative number modify the
          brightness by this factor.

    Returns:
        torch.Tensor: Adjusted image.
    """

    def __init__(self, contrast_factor: Union[float, torch.Tensor]) -> None:
        super(AdjustContrast, self).__init__()
        self.contrast_factor: Union[float, torch.Tensor] = contrast_factor

    def forward(self, input: torch.Tensor) -> torch.Tensor:  # type: ignore
        return adjust_contrast(input, self.contrast_factor)


class AdjustBrightness(nn.Module):
    r"""Adjust Brightness of an image. This implementation aligns OpenCV, not PIL. Hence,
    the output differs from TorchVision.

    The input image is expected to be in the range of [0, 1].

    Args:
        input (torch.Tensor): Image/Input to be adjusted in the shape of (\*, N).
        brightness_factor (Union[float, torch.Tensor]): Brightness adjust factor per element
          in the batch. 0 does not modify the input image while any other number modify the
          brightness.

    Returns:
        torch.Tensor: Adjusted image.
    """

    def __init__(self, brightness_factor: Union[float, torch.Tensor]) -> None:
        super(AdjustBrightness, self).__init__()
        self.brightness_factor: Union[float, torch.Tensor] = brightness_factor

    def forward(self, input: torch.Tensor) -> torch.Tensor:  # type: ignore
        return adjust_brightness(input, self.brightness_factor)

[((1212, 1248), 'torch.chunk', 'torch.chunk', (['input'], {'chunks': '(3)', 'dim': '(-3)'}), '(input, chunks=3, dim=-3)\n', (1223, 1248), False, 'import torch\n'), ((1322, 1370), 'torch.clamp', 'torch.clamp', (['(s * saturation_factor)'], {'min': '(0)', 'max': '(1)'}), '(s * saturation_factor, min=0, max=1)\n', (1333, 1370), False, 'import torch\n'), ((1435, 1467), 'torch.cat', 'torch.cat', (['[h, s_out, v]'], {'dim': '(-3)'}), '([h, s_out, v], dim=-3)\n', (1444, 1467), False, 'import torch\n'), ((1770, 1787), 'kornia.color.hsv.rgb_to_hsv', 'rgb_to_hsv', (['input'], {}), '(input)\n', (1780, 1787), False, 'from kornia.color.hsv import rgb_to_hsv, hsv_to_rgb\n'), ((1948, 1970), 'kornia.color.hsv.hsv_to_rgb', 'hsv_to_rgb', (['x_adjusted'], {}), '(x_adjusted)\n', (1958, 1970), False, 'from kornia.color.hsv import rgb_to_hsv, hsv_to_rgb\n'), ((2993, 3029), 'torch.chunk', 'torch.chunk', (['input'], {'chunks': '(3)', 'dim': '(-3)'}), '(input, chunks=3, dim=-3)\n', (3004, 3029), False, 'import torch\n'), ((3138, 3173), 'torch.fmod', 'torch.fmod', (['(h + hue_factor)', 'divisor'], {}), '(h + hue_factor, divisor)\n', (3148, 3173), False, 'import torch\n'), ((3238, 3270), 'torch.cat', 'torch.cat', (['[h_out, s, v]'], {'dim': '(-3)'}), '([h_out, s, v], dim=-3)\n', (3247, 3270), False, 'import torch\n'), ((3539, 3556), 'kornia.color.hsv.rgb_to_hsv', 'rgb_to_hsv', (['input'], {}), '(input)\n', (3549, 3556), False, 'from kornia.color.hsv import rgb_to_hsv, hsv_to_rgb\n'), ((3703, 3725), 'kornia.color.hsv.hsv_to_rgb', 'hsv_to_rgb', (['x_adjusted'], {}), '(x_adjusted)\n', (3713, 3725), False, 'from kornia.color.hsv import rgb_to_hsv, hsv_to_rgb\n'), ((5135, 5166), 'torch.clamp', 'torch.clamp', (['x_adjust', '(0.0)', '(1.0)'], {}), '(x_adjust, 0.0, 1.0)\n', (5146, 5166), False, 'import torch\n'), ((6288, 6319), 'torch.clamp', 'torch.clamp', (['x_adjust', '(0.0)', '(1.0)'], {}), '(x_adjust, 0.0, 1.0)\n', (6299, 6319), False, 'import torch\n'), ((7347, 7378), 'torch.clamp', 'torch.clamp', (['x_adjust', '(0.0)', '(1.0)'], {}), '(x_adjust, 0.0, 1.0)\n', (7358, 7378), False, 'import torch\n'), ((431, 453), 'torch.is_tensor', 'torch.is_tensor', (['input'], {}), '(input)\n', (446, 453), False, 'import torch\n'), ((825, 858), 'torch.tensor', 'torch.tensor', (['[saturation_factor]'], {}), '([saturation_factor])\n', (837, 858), False, 'import torch\n'), ((1126, 1168), 'torch.unsqueeze', 'torch.unsqueeze', (['saturation_factor'], {'dim': '(-1)'}), '(saturation_factor, dim=-1)\n', (1141, 1168), False, 'import torch\n'), ((2238, 2260), 'torch.is_tensor', 'torch.is_tensor', (['input'], {}), '(input)\n', (2253, 2260), False, 'import torch\n'), ((2628, 2654), 'torch.tensor', 'torch.tensor', (['[hue_factor]'], {}), '([hue_factor])\n', (2640, 2654), False, 'import torch\n'), ((2914, 2949), 'torch.unsqueeze', 'torch.unsqueeze', (['hue_factor'], {'dim': '(-1)'}), '(hue_factor, dim=-1)\n', (2929, 2949), False, 'import torch\n'), ((3102, 3111), 'kornia.constants.pi.item', 'pi.item', ([], {}), '()\n', (3109, 3111), False, 'from kornia.constants import pi\n'), ((4013, 4035), 'torch.is_tensor', 'torch.is_tensor', (['input'], {}), '(input)\n', (4028, 4035), False, 'import torch\n'), ((4472, 4493), 'torch.tensor', 'torch.tensor', (['[gamma]'], {}), '([gamma])\n', (4484, 4493), False, 'import torch\n'), ((4542, 4562), 'torch.tensor', 'torch.tensor', (['[gain]'], {}), '([gain])\n', (4554, 4562), False, 'import torch\n'), ((4904, 4934), 'torch.unsqueeze', 'torch.unsqueeze', (['gamma'], {'dim': '(-1)'}), '(gamma, dim=-1)\n', (4919, 4934), False, 'import torch\n'), ((4950, 4979), 'torch.unsqueeze', 'torch.unsqueeze', (['gain'], {'dim': '(-1)'}), '(gain, dim=-1)\n', (4965, 4979), False, 'import torch\n'), ((5050, 5073), 'torch.pow', 'torch.pow', (['input', 'gamma'], {}), '(input, gamma)\n', (5059, 5073), False, 'import torch\n'), ((5425, 5447), 'torch.is_tensor', 'torch.is_tensor', (['input'], {}), '(input)\n', (5440, 5447), False, 'import torch\n'), ((5801, 5832), 'torch.tensor', 'torch.tensor', (['[contrast_factor]'], {}), '([contrast_factor])\n', (5813, 5832), False, 'import torch\n'), ((6088, 6128), 'torch.unsqueeze', 'torch.unsqueeze', (['contrast_factor'], {'dim': '(-1)'}), '(contrast_factor, dim=-1)\n', (6103, 6128), False, 'import torch\n'), ((6588, 6610), 'torch.is_tensor', 'torch.is_tensor', (['input'], {}), '(input)\n', (6603, 6610), False, 'import torch\n'), ((6972, 7005), 'torch.tensor', 'torch.tensor', (['[brightness_factor]'], {}), '([brightness_factor])\n', (6984, 7005), False, 'import torch\n'), ((7141, 7183), 'torch.unsqueeze', 'torch.unsqueeze', (['brightness_factor'], {'dim': '(-1)'}), '(brightness_factor, dim=-1)\n', (7156, 7183), False, 'import torch\n')]