# Copyright 2004-present Facebook. All Rights Reserved.
#
# Author: Daniel DeTone (ddetone)

import math
import warnings
from typing import NamedTuple, Optional

import torch


def truncated_normal_(N: int, mean: float = 0.0, std: float = 1):
    """Draws N samples from a truncated normal distribution ~N_t(mean,std).

    :param N: number of samples to draw
    :param mean: mean of distribution
    :param std: standard deviation of distribution

    """
    tensor = torch.FloatTensor(N)
    tmp = tensor.new_empty(tensor.shape + (4,)).normal_()
    valid = (tmp < 2) & (tmp > -2)
    ind = valid.max(-1, keepdim=True)[1]
    tensor.data.copy_(tmp.gather(-1, ind).squeeze(-1))
    tensor.data.mul_(std).add_(mean)
    return tensor


def perspective_grid(
    coeffs: torch.Tensor,
    size: torch.Size,
    align_corners: bool = False,
    eps: float = 1e-9,
):
    """
    Generate a grid of u,v sample locations from a perspective transform.
    Acts similarly to torch's affine_grid, but allows for 3x3 perspective warps.

    :param coeffs: Bx8 input batch of perspective transform matrices
    (flattened first 8 elements of 3x3 matrix, assuming last element = 1)
    :param size: NxCxHxW the target output size
    :param align_corners: if True, consider -1 and 1 to refer to the centers
    of the corner pixels
    :param eps: factor to prevent division by zero

    """
    B, H, W = int(size[0]), int(size[-2]), int(size[-1])
    base_grid = torch.empty(B, H, W, 3).to(coeffs)
    # Generate base_grid [-1,1], taking into account align_corners.
    step_x = torch.linspace(-1, 1, W)
    step_y = torch.linspace(-1, 1, H)
    if align_corners is False:
        step_x = step_x * (W - 1) / W
        step_y = step_y * (H - 1) / H
    base_grid[..., 0].copy_(step_x)
    base_grid[..., 1].copy_(step_y.unsqueeze(-1))
    base_grid[..., 2].fill_(1)
    base_grid = base_grid.reshape(B, -1, 3)
    # Generate homography matrix.
    h_mat = torch.ones(B, 9).to(coeffs)
    h_mat[:, :-1] = coeffs
    h_mat = h_mat.reshape(B, 3, 3).transpose(1, 2)
    # Apply homography warp, then normalize by homogenous coordinate.
    grid = base_grid @ h_mat
    grid[:, :, 0] = grid[:, :, 0] / (grid[:, :, 2] + eps)
    grid[:, :, 1] = grid[:, :, 1] / (grid[:, :, 2] + eps)
    grid = grid[:, :, :2].reshape(B, H, W, 2)
    return grid


def check_input(inp: torch.Tensor):
    """Makes sure that input tensor meets input specifications."""
    if inp.dim() != 4:
        raise ValueError("Input tensor must have 4 dims (B, C, H, W)")
    if inp.shape[1] != 1 and inp.shape[1] != 3:
        raise ValueError("Only 1 and 3 channel inputs are supported")
    if inp.dtype != torch.float32:
        raise TypeError("Only torch.float32 supported.")
    if inp.min() < 0.0 or inp.max() > 1.0:
        raise ValueError("Image must be normalized between [0,1].")
    return


# Note(dd): It was a pain to update kornia, so I took this from the kornia directly.
def get_perspective_transform(src, dst):
    r"""Calculates a perspective transform from four pairs of the corresponding
    points.
    The function calculates the matrix of a perspective transform so that:
    .. math ::
        \begin{bmatrix}
        t_{i}x_{i}^{'} \\
        t_{i}y_{i}^{'} \\
        t_{i} \\
        \end{bmatrix}
        =
        \textbf{map_matrix} \cdot
        \begin{bmatrix}
        x_{i} \\
        y_{i} \\
        1 \\
        \end{bmatrix}
    where
    .. math ::
        dst(i) = (x_{i}^{'},y_{i}^{'}), src(i) = (x_{i}, y_{i}), i = 0,1,2,3
    Args:
        src (torch.Tensor): coordinates of quadrangle vertices in the source image
        with shape :math:`(B, 4, 2)`.
        dst (torch.Tensor): coordinates of the corresponding quadrangle vertices in
            the destination image with shape :math:`(B, 4, 2)`.
    Returns:
        torch.Tensor: the perspective transformation with shape :math:`(B, 3, 3)`.
    .. note::
        This function is often used in conjuntion with :func:`warp_perspective`.
    """

    # we build matrix A by using only 4 point correspondence. The linear
    # system is solved with the least square method, so here
    # we could even pass more correspondence
    p = []
    for i in [0, 1, 2, 3]:
        p.append(_build_perspective_param(src[:, i], dst[:, i], "x"))
        p.append(_build_perspective_param(src[:, i], dst[:, i], "y"))

    # A is Bx8x8
    A = torch.stack(p, dim=1)

    # b is a Bx8x1
    b = torch.stack(
        [
            dst[:, 0:1, 0],
            dst[:, 0:1, 1],
            dst[:, 1:2, 0],
            dst[:, 1:2, 1],
            dst[:, 2:3, 0],
            dst[:, 2:3, 1],
            dst[:, 3:4, 0],
            dst[:, 3:4, 1],
        ],
        dim=1,
    )

    # solve the system Ax = b
    X = torch.linalg.solve(A, b)

    # create variable to return
    batch_size = src.shape[0]
    M = torch.ones(batch_size, 9, device=src.device, dtype=src.dtype)
    M[..., :8] = torch.squeeze(X, dim=-1)

    return M.view(-1, 3, 3)  # Bx3x3


def _build_perspective_param(
    p: torch.Tensor, q: torch.Tensor, axis: str
) -> torch.Tensor:
    ones = torch.ones_like(p)[..., 0:1]
    zeros = torch.zeros_like(p)[..., 0:1]
    if axis == "x":
        return torch.cat(
            [
                p[:, 0:1],
                p[:, 1:2],
                ones,
                zeros,
                zeros,
                zeros,
                -p[:, 0:1] * q[:, 0:1],
                -p[:, 1:2] * q[:, 0:1],
            ],
            dim=1,
        )

    elif axis == "y":
        return torch.cat(
            [
                zeros,
                zeros,
                zeros,
                p[:, 0:1],
                p[:, 1:2],
                ones,
                -p[:, 0:1] * q[:, 1:2],
                -p[:, 1:2] * q[:, 1:2],
            ],
            dim=1,
        )
    else:
        raise ValueError("Bad input axis, should be x or y")


class GeoAugParam(NamedTuple):
    min: Optional[float] = 0.0
    max: Optional[float] = 0.0


class RandomGeoAug(object):
    """
    Applies batch-wise geometric augmentations to images.
    Requires that input/output tensors are sized BxCxHxW.
    """

    def __init__(
        self,
        rotate_param: Optional[GeoAugParam] = None,
        scale_param: Optional[GeoAugParam] = None,
        translate_x_param: Optional[GeoAugParam] = None,
        translate_y_param: Optional[GeoAugParam] = None,
        shear_x_param: Optional[GeoAugParam] = None,
        shear_y_param: Optional[GeoAugParam] = None,
        perspective_param: Optional[GeoAugParam] = None,
    ):

        self.rotate_param = (
            GeoAugParam(min=-180, max=180) if rotate_param is None else rotate_param
        )
        self.scale_param = (
            GeoAugParam(min=0.5, max=2.0) if scale_param is None else scale_param
        )
        self.translate_x_param = (
            GeoAugParam(min=-0.3, max=0.3)
            if translate_x_param is None
            else translate_x_param
        )
        self.translate_y_param = (
            GeoAugParam(min=-0.3, max=0.3)
            if translate_y_param is None
            else translate_y_param
        )
        self.shear_x_param = (
            GeoAugParam(min=-20, max=20) if shear_x_param is None else shear_x_param
        )
        self.shear_y_param = (
            GeoAugParam(min=-20, max=20) if shear_y_param is None else shear_y_param
        )
        self.perspective_param = (
            GeoAugParam(min=-0.3, max=0.3)
            if perspective_param is None
            else perspective_param
        )
        self.eps = 1e-9

        self.param_names = []
        for key in dir(self):
            if "_param" in key:
                self.param_names.append(key)

    def __repr__(self):
        ret = self.__class__.__name__ + "(\n"
        for key in dir(self):
            if "_param" in key:
                ret += ("  " + key + " = " + getattr(self, key).__repr__()) + ",\n"
        ret += ")"
        return ret

    def _get_perspective_matrix(
        self,
        angle: torch.Tensor,
        translate: torch.Tensor,
        scale: torch.Tensor,
        shear: torch.Tensor,
        delta_corners: torch.Tensor,
    ):
        B = angle.shape[0]
        # Thus, the inverse is M^-1 = RSS^-1 * T^-1
        rot = torch.deg2rad(-angle)
        sx, sy = [torch.deg2rad(s) for s in shear.t()]
        tx, ty = 2 * translate[:, 0], 2 * translate[:, 1]
        # RSS without scaling
        a = torch.cos(rot - sy) / torch.cos(sy)
        b = -torch.cos(rot - sy) * torch.tan(sx) / torch.cos(sy) - torch.sin(rot)
        c = torch.sin(rot - sy) / torch.cos(sy)
        d = -torch.sin(rot - sy) * torch.tan(sx) / torch.cos(sy) + torch.cos(rot)
        # Inverted rotation matrix with scale and shear
        zeros = torch.zeros_like(a)
        matrix = torch.stack([d, -b, zeros, -c, a, zeros], dim=-1)
        matrix = torch.stack([x / scale for x in matrix.t()], dim=-1)
        # Apply inverse of translation and of center translation: RSS^-1 * T^-1
        matrix[:, 2] += matrix[:, 0] * (-tx) + matrix[:, 1] * (-ty)
        matrix[:, 5] += matrix[:, 3] * (-tx) + matrix[:, 4] * (-ty)
        matrix = torch.cat([matrix, matrix.new_zeros((B, 3))], dim=1)
        matrix[:, -1] = 1.0
        matrix = matrix.reshape(B, 3, 3)
        # Generate non-affine effect by perturbing the four corners.
        start = torch.ones(B, 4, 2).to(a)
        end = start - delta_corners  # Positive delta takes corners inwards.
        corner_signs = torch.tensor([[-1, -1], [1, -1], [1, 1], [-1, 1]]).to(a)
        corner_signs = corner_signs.repeat(B, 1, 1)
        start *= corner_signs
        end *= corner_signs
        # Set up a linear system to solve for the homography.
        # pmat = solve_homography(start, end)
        # pmat = solve_homography_old(start, end)
        # pmat = kornia.homography.find_homography_dlt(start, end) # Newer kornia version API.
        # pmat = kornia.get_perspective_transform(start, end)
        pmat = get_perspective_transform(
            start, end
        )  # Copied kornia functioned directly.
        # Apply perspective transform first, then affine.
        matrix = pmat @ matrix
        # Normalize such that H[2,2] = 1.
        eps = 1e-9
        matrix /= matrix[:, 2, 2].clone().reshape(B, 1, 1) + eps
        return matrix

    def _generate_transform(self, B: int, device: torch.device):
        rot_range = (self.rotate_param.min, self.rotate_param.max)
        trans_x_range = (self.translate_x_param.min, self.translate_x_param.max)
        trans_y_range = (self.translate_y_param.min, self.translate_y_param.max)
        shear_x_range = (self.shear_x_param.min, self.shear_x_param.max)
        shear_y_range = (self.shear_y_param.min, self.shear_y_param.max)
        scale_range = (self.scale_param.min, self.scale_param.max)
        pscale_range = (self.perspective_param.min, self.perspective_param.max)
        # Randomly generate parameters that define the perspective transform.
        rot = torch.empty(B, device=device).uniform_(rot_range[0], rot_range[1])
        trans_x = torch.empty(B, device=device).uniform_(
            trans_x_range[0], trans_x_range[1]
        )
        trans_y = torch.empty(B, device=device).uniform_(
            trans_y_range[0], trans_y_range[1]
        )
        trans_xy = torch.stack([trans_x, trans_y], dim=-1)
        shear_x = torch.empty(B, device=device).uniform_(
            shear_x_range[0], shear_x_range[1]
        )
        shear_y = torch.empty(B, device=device).uniform_(
            shear_y_range[0], shear_y_range[1]
        )
        shear_xy = torch.stack([shear_x, shear_y], dim=-1)
        scale = torch.empty(B, device=device).uniform_(scale_range[0], scale_range[1])
        delta_corners = torch.empty(B, 4, 2, device=device).uniform_(
            pscale_range[0], pscale_range[1]
        )
        # Given parameters, generate 3x3 perspective matrices.
        transform = self._get_perspective_matrix(
            rot, trans_xy, scale, shear_xy, delta_corners
        )
        return transform

    def _apply_perspective(
        self, img: torch.Tensor, transform: torch.Tensor, align_corners: bool = True
    ):
        coeffs = transform.reshape(-1, 9)[:, :-1]
        grid2 = perspective_grid(
            coeffs, img.shape, align_corners=align_corners, eps=self.eps
        )
        img2 = torch.nn.functional.grid_sample(img, grid2, align_corners=align_corners)
        return img2

    def _denormalize_transform(self, transform: torch.Tensor, H: int, W: int):
        B = int(transform.shape[0])
        ones = transform.new_ones((B, 1))
        zeros = transform.new_zeros((B, 1))
        scaleW = (1.0 / ((W - 1.0) / 2.0)) * ones
        scaleH = (1.0 / ((H - 1.0) / 2.0)) * ones
        T = torch.stack(
            [scaleW, zeros, -ones, zeros, scaleH, -ones, zeros, zeros, ones], dim=-1
        ).reshape(-1, 3, 3)
        transform2 = torch.inverse(T) @ transform @ T
        transform2 /= transform2[:, 2, 2].reshape(B, 1, 1).clone() + self.eps
        return transform2

    def set_all_identity(self):
        for name in self.param_names:
            if "scale_param" in name:
                val = 1.0
            else:
                val = 0.0
            setattr(self, name, GeoAugParam(min=val, max=val))

    def forward(self, inp, return_transform=False, normalize_returned_transform=False):
        """
        Runs the Random Geometric Augmentation on a batch of images.

        :param inp: BxCxHxW batch of input images
        :param return_transform: if True, also return the transform used to
        warp the inp (in unnormalize image coordinates by default)
        :param normalize_returned_transform: if True, the returned transform
        will be in normalized pixel coords (e.g. [-1,-1] is top left pixel,
        [1,1] is bottom right)

        """
        check_input(inp)  # Make sure input is well-formed.
        if self.scale_param.min <= 0.0 or self.scale_param.max <= 0.0:
            warnings.warn(
                "Scale min and max are are too extreme, resulting image might "
                "be black, range should be (0, inf)",
                RuntimeWarning,
            )
        if self.translate_x_param.min < -1.0 or self.translate_x_param.max > 1.0:
            warnings.warn(
                "Translate_x min and max are too extreme, resulting image might"
                "be black, range should be [-1.0, 1.0]",
                RuntimeWarning,
            )
        if self.translate_y_param.min < -1.0 or self.translate_y_param.max > 1.0:
            warnings.warn(
                "Translate_y min and max are too extreme, resulting image "
                "might be black, range should be [-1.0, 1.0]",
                RuntimeWarning,
            )
        if self.shear_x_param.min < -90 or self.shear_x_param.max > 90:
            warnings.warn(
                "Shear_x min and max are too extreme, resulting image might "
                "be black, range should be [-90, 90]",
                RuntimeWarning,
            )
        if self.shear_y_param.min < -90 or self.shear_y_param.max > 90:
            warnings.warn(
                "Shear_x min and max are too extreme, resulting image might "
                "be black, range should be [-90, 90]",
                RuntimeWarning,
            )
        if self.perspective_param.min < -1.0 or self.perspective_param.max < -1.0:
            warnings.warn(
                "Perspective min and max are too extreme, resulting image "
                "might be black, range should be [-1, inf]",
                RuntimeWarning,
            )

        # Forcing this to be generated on CPU because it's faster.
        inv_transform = self._generate_transform(
            inp.shape[0], device=torch.device("cpu")
        )
        inv_transform = inv_transform.to(inp)
        identity = (
            torch.eye(3).unsqueeze(0).repeat(inp.shape[0], 1, 1).to(inv_transform)
        )
        if torch.allclose(identity, inv_transform):  # Shortcut operation.
            inp2 = inp
            transform = inv_transform
        else:
            inp2 = self._apply_perspective(inp, inv_transform, align_corners=True)
            transform = torch.inverse(inv_transform)
            transform /= transform[:, 2, 2].reshape(-1, 1, 1).clone() + self.eps
        if not return_transform:
            return inp2
        else:
            if normalize_returned_transform:
                return inp2, transform
            else:
                H, W = inp.shape[-2], inp.shape[-1]
                denorm_transform = self._denormalize_transform(transform, H, W)
                return inp2, denorm_transform


class PhotoAugParam(NamedTuple):
    prob: Optional[float] = 0.2
    mag: Optional[float] = 0.5
    fix_pos: Optional[bool] = False


class RandomPhotoAug(object):
    """
    Applies batch-wise photometric distortions to images.
    Requires that input/output tensors are sized BxCxHxW.

    """

    def __init__(
        self,
        contrast_param: Optional[PhotoAugParam] = None,
        sharpen_param: Optional[PhotoAugParam] = None,
        exposure_param: Optional[PhotoAugParam] = None,
        gamma_param: Optional[PhotoAugParam] = None,
        gaussian_smooth_param: Optional[PhotoAugParam] = None,
        motion_blur_param: Optional[PhotoAugParam] = None,
        shadow_highlight_param: Optional[PhotoAugParam] = None,
        gaussian_noise_param: Optional[PhotoAugParam] = None,
        salt_and_pepper_param: Optional[PhotoAugParam] = None,
    ):

        self.contrast_param = (
            PhotoAugParam() if contrast_param is None else contrast_param
        )
        self.sharpen_param = PhotoAugParam() if sharpen_param is None else sharpen_param
        self.exposure_param = (
            PhotoAugParam() if exposure_param is None else exposure_param
        )
        self.gamma_param = PhotoAugParam() if gamma_param is None else gamma_param
        self.gaussian_smooth_param = (
            PhotoAugParam() if gaussian_smooth_param is None else gaussian_smooth_param
        )
        self.motion_blur_param = (
            PhotoAugParam() if motion_blur_param is None else motion_blur_param
        )
        self.shadow_highlight_param = (
            PhotoAugParam()
            if shadow_highlight_param is None
            else shadow_highlight_param
        )
        # Fix the order in which the last two are applied by default.
        self.gaussian_noise_param = (
            PhotoAugParam(fix_pos=True)
            if gaussian_noise_param is None
            else gaussian_noise_param
        )
        self.salt_and_pepper_param = (
            PhotoAugParam(fix_pos=True)
            if salt_and_pepper_param is None
            else salt_and_pepper_param
        )

        # Sets the default order.
        self.fn_names = [
            "contrast",
            "sharpen",
            "exposure",
            "gamma",
            "gaussian_smooth",
            "motion_blur",
            "shadow_highlight",
            "gaussian_noise",
            "salt_and_pepper",
        ]

        # Each fn_name must be a method.
        for name in self.fn_names:
            assert hasattr(self, name) and callable(getattr(self, name))

        self.param_names = []
        for key in dir(self):
            if "_param" in key:
                self.param_names.append(key)

    def __repr__(self):
        ret = self.__class__.__name__ + "(\n"
        for key in dir(self):
            if key.endswith("_param"):
                ret += ("  " + key + " = " + getattr(self, key).__repr__()) + ",\n"
        ret += ")"
        return ret

    def set_all_probs(self, val: int):
        for name in self.param_names:
            param = getattr(self, name)
            setattr(self, name, param._replace(prob=val))

    def set_all_mags(self, val: int):
        for name in self.param_names:
            param = getattr(self, name)
            setattr(self, name, param._replace(mag=val))

    def set_all_fix_pos(self, val: bool):
        for name in self.param_names:
            param = getattr(self, name)
            setattr(self, name, param._replace(fix_pos=val))

    def contrast(self, inp: torch.Tensor, mag: float = 0.5):
        B = inp.shape[0]
        # Kind of arbitrary, asymmetric normal distribution with
        # has median=1 and doesn't have clipping artifacts.
        scales = truncated_normal_(B, mean=1.0, std=(mag * 0.9 * 0.5 + 1e-9))
        scales = torch.where(scales > 1.0, scales + (scales - 1.0) * 3.0, scales)
        scales = scales.to(inp)
        out = inp - 0.5
        out = out * scales.reshape(-1, 1, 1, 1)
        out = out + 0.5
        out = torch.clamp(out, 0.0, 1.0)
        return out

    def sharpen(self, inp: torch.Tensor, mag: float = 0.5):
        B, C, H, W = inp.shape
        scale = truncated_normal_(B, mean=0, std=mag).abs()
        val = (1.0 / 5.0) + scale
        kernel = torch.zeros(B, 3, 3).to(inp)
        kernel[:, 1, 1] = val * 5
        kernel[:, 0, 1] = -((5 * val - 1.0) / 4.0)
        kernel[:, 1, 0] = -((5 * val - 1.0) / 4.0)
        kernel[:, 1, 2] = -((5 * val - 1.0) / 4.0)
        kernel[:, 2, 1] = -((5 * val - 1.0) / 4.0)
        # Use grouped conv to apply different kernel to each batch element.
        kernel = kernel[:, None, ...]
        kernel = torch.repeat_interleave(kernel, C, dim=0)
        inp = inp.reshape(1, B * C, H, W)
        out = torch.nn.functional.conv2d(inp, kernel, padding=1, groups=B * C)
        out = out.reshape(B, C, H, W)
        out = torch.clamp(out, 0.0, 1.0)
        return out

    def exposure(self, inp: torch.Tensor, mag: float = 0.5):
        B = inp.shape[0]
        max_delta = mag * 0.9  # Don't let max pixel value go below 0.1.
        new_maxs = truncated_normal_(B, mean=1.0, std=0.5 * max_delta).to(inp)
        out = inp * new_maxs.reshape(-1, 1, 1, 1)
        out = torch.clamp(out, 0.0, 1.0)
        return out

    def gamma(self, inp: torch.Tensor, mag: float = 0.5):
        B = inp.shape[0]
        # Kind of arbitrary, asymmetric normal distribution with
        # has median=1 and doesn't have clipping artifacts.
        gam = truncated_normal_(B, mean=1.0, std=mag * 0.9 * 0.5).to(inp)
        gam = torch.where(gam > 1.0, gam + (gam - 1.0) * 3.0, gam)
        out = (inp ** gam.reshape(-1, 1, 1, 1)).clamp(0, 1)
        return out

    def gaussian_smooth(
        self,
        inp: torch.Tensor,
        kmin: int = 1,
        kmax: int = 11,
        smin: float = 1.0,
        smax: float = 5.0,
        mag: float = 0.5,
    ):
        B, C, H, W = inp.shape
        # Maximum possible kernel size.
        kmax = int(kmax * mag) + 1
        if kmax == 1:  # 1x1 gaussian blur does nothing.
            return inp
        half_sizes = torch.randint(size=(B,), low=kmin, high=kmax)  # check kmax
        # To batch the operations, we use the largest kernel size.
        kernel_size = int(half_sizes.max().item() * 2 + 1)
        sigma = mag * torch.FloatTensor(B).uniform_(smin, smax)
        x_cord = torch.arange(kernel_size)
        x_grid = x_cord.repeat(kernel_size).view(kernel_size, kernel_size)
        y_grid = x_grid.t()
        xy_grid = torch.stack([x_grid, y_grid], dim=-1)
        xy_grid = xy_grid.unsqueeze(0).repeat(B, 1, 1, 1)
        mean = (kernel_size - 1) / 2.0
        variance = sigma**2.0
        variance = variance.reshape(-1, 1, 1, 1)
        gaussian_kernel = (1.0 / (2.0 * math.pi * variance)) * torch.exp(
            -torch.sum((xy_grid - mean) ** 2.0, dim=-1, keepdim=True) / (2 * variance)
        )
        # Make sure sum of values in gaussian kernel equals 1.
        kernel = gaussian_kernel / torch.sum(
            gaussian_kernel, dim=[1, 2, 3], keepdim=True
        )
        kernel = kernel.to(inp)
        # Use grouped conv to apply different kernel to each batch element.
        kernel = kernel.reshape(B, 1, kernel_size, kernel_size)
        kernel = torch.repeat_interleave(kernel, C, dim=0)
        kernel = kernel.reshape(B * C, 1, kernel_size, kernel_size)
        inp = inp.reshape(1, B * C, H, W)
        out = torch.nn.functional.conv2d(
            inp, kernel, padding=int(half_sizes.max()), groups=B * C
        )
        out = out.reshape(B, C, H, W)
        out = torch.clamp(out, 0.0, 1.0)
        return out

    def motion_blur(
        self, inp: torch.Tensor, mag: float = 0.5, kmin: int = 1, kmax: int = 11
    ):
        B, C, H, W = inp.shape
        # Maximum possible kernel size.
        kmax = int(kmax * mag) + 1
        if kmax == 1:  # 1x1 gaussian blur does nothing.
            return inp
        half_sizes = torch.randint(size=(B,), low=kmin, high=kmax)  # check kmax
        # To batch the operations, we use the largest kernel size.
        kernel_sizes = half_sizes * 2 + 1
        kernel_size = int(kernel_sizes.max().item())
        # Generate a different motion blur kernel for each batch.
        kernel = torch.zeros((B, kernel_size, kernel_size)).to(inp)
        # TODO(dd): figure out how to vectorize this.
        for i, ks in enumerate(kernel_sizes):
            off = torch.div((kernel_size - ks), 2, rounding_mode="floor")
            start = int(off)
            end = int(kernel_size - off)
            kh = torch.div(kernel_size, 2, rounding_mode="floor")
            kernel[i, kh, start:end] = torch.ones((1, ks)) / ks
        kernel = kernel[:, None, :, :]
        # Rotate the motion blur kernel.
        angle = torch.FloatTensor(B).uniform_(0, 360.0).unsqueeze(0).to(inp)
        rads = (angle * (math.pi / 180.0)).reshape(B, 1, 1)
        r00 = torch.cos(rads)
        r01 = -torch.sin(rads)
        r10 = torch.sin(rads)
        r11 = torch.cos(rads)
        top = torch.cat([r00, r01, 0 * r00], dim=2)
        bot = torch.cat([r10, r11, 0 * r00], dim=2)
        affine = torch.cat([top, bot], dim=1).to(inp)
        grid = torch.nn.functional.affine_grid(
            affine, list(kernel.shape), align_corners=True
        )
        kernel = torch.nn.functional.grid_sample(kernel, grid, align_corners=True)
        # Apply the rotated motion blur kernel to the images.
        kernel = torch.repeat_interleave(kernel, C, dim=0)
        kernel = kernel.reshape(B * C, 1, kernel_size, kernel_size)
        inp = inp.reshape(1, B * C, H, W)
        out = torch.nn.functional.conv2d(
            inp, kernel, padding=int(half_sizes.max()), groups=B * C
        )
        out = out.reshape(B, C, H, W)
        out = torch.clamp(out, 0.0, 1.0)
        return out

    def shadow_highlight(
        self,
        inp: torch.Tensor,
        mag: float = 0.5,
        min_quads: int = 2,
        max_quads: int = 5,
    ):
        B, H, W = inp.shape[0], inp.shape[-2], inp.shape[-1]
        out = inp
        num_quads = torch.randint(low=min_quads, high=max_quads, size=(B,))

        kwargs = {
            "scale_param": GeoAugParam(min=0.5, max=0.9),
            "translate_x_param": GeoAugParam(min=-0.8, max=0.8),
            "translate_y_param": GeoAugParam(min=-0.8, max=0.8),
            "rotate_param": GeoAugParam(min=-180, max=180),
            "perspective_param": GeoAugParam(min=0.0, max=0.8),
        }
        rga = RandomGeoAug(**kwargs)
        max_num_quads = num_quads.max()
        for i in range(max_num_quads):
            # Faster to apply conv to mini image then resize later.
            mask = torch.ones(B, 1, 50, 50).to(inp)
            # Randomly distort the quad.
            mask = rga.forward(mask)
            # Smooth the quad.
            mask = self.gaussian_smooth(mask, kmin=3, kmax=11, smin=3, smax=11, mag=1.0)
            mask = torch.nn.functional.interpolate(
                mask, size=[H, W], mode="bilinear", align_corners=True
            )
            val = truncated_normal_(B, mean=0.0, std=(mag * 0.6)).to(inp)
            quad = mask * val.reshape(B, 1, 1, 1)
            out2 = (1.0 - mask) * out + mask * (out + quad)
            apply_quad = torch.tensor(i).reshape(1, 1, 1, 1).repeat(
                B, 1, 1, 1
            ) < num_quads.reshape(B, 1, 1, 1)
            apply_quad = apply_quad.to(inp.device)
            out = torch.where(apply_quad, out2, out)
            out = torch.clamp(out, 0.0, 1.0)
        return out

    def gaussian_noise(self, inp: torch.Tensor, mag: float = 0.5):
        B, C, H, W = inp.shape
        val = torch.FloatTensor(B).uniform_(0, mag * 0.1).abs().reshape(B, 1, 1)
        noise = val * torch.randn((B, H, W))
        noise = noise[:, None, :, :]
        noise = noise.repeat(1, C, 1, 1).to(inp.device)
        out = inp + noise
        out = torch.clamp(out, 0.0, 1.0)
        return out

    def salt_and_pepper(self, inp: torch.Tensor, mag: float = 0.5):
        B, C, H, W = inp.shape
        noise_prob = mag * torch.FloatTensor(B).normal_(
            0.0, (mag * 0.05 + 1e-9)
        ).abs().to(inp)
        coins = torch.rand(size=(B, H, W)).to(inp).unsqueeze(1).repeat(1, C, 1, 1)
        heavy_noise = 0.5 * torch.randn(B, H, W).to(inp).unsqueeze(1).repeat(1, C, 1, 1)
        out = torch.where(
            coins > noise_prob.reshape(-1, 1, 1, 1), inp, inp + heavy_noise
        )
        out = torch.clamp(out, 0.0, 1.0)
        return out

    def forward(self, inp: torch.Tensor, debug: bool = False, seed: int = None):
        """
        Runs the Random Photometric Augmentation on a batch of images.

        :param inp: BxCxHxW batch of input images
        :param debug: if True, print out some info about which transforms were used

        """
        if seed is not None:
            torch.manual_seed(seed)
        check_input(inp)  # Make sure input is well-formed.
        B, C, H, W = inp.shape
        result = inp.clone()  # Creates a copy of input (leaves input unchanged).
        # Randomly permute function ordering, except if "fix_pos" is True.
        permute_names = [
            name for name in self.fn_names if not getattr(self, name + "_param").fix_pos
        ]
        indices = torch.randperm(len(permute_names)).tolist()
        perm_fns = [permute_names[idx] for idx in indices]
        for i, name in enumerate(self.fn_names):
            if getattr(self, name + "_param").fix_pos:
                perm_fns.insert(i, name)
        # Iterate randomly permuted augmentation functions and maybe apply it.
        for fn_name in perm_fns:
            prob = getattr(self, fn_name + "_param").prob
            coins = torch.rand(B).reshape(-1, 1, 1, 1).to(inp)
            mask = coins < prob
            # Optional printing, helpful for debugging.
            if debug:
                print(fn_name, int(mask.sum()), "/", result.shape[0])
            # Skip operation if no elt in batch needs this augmentation.
            if torch.all(torch.eq(mask, False)):
                continue
            # Run the augmentation across all elts in batch.
            mag = getattr(self, fn_name + "_param").mag
            fn_method = getattr(self, fn_name)
            pre = result.clone()
            # Actual call to augmentation is here.
            result2 = fn_method(pre, mag=mag)
            mask = mask.repeat(1, C, H, W)
            # Apply the augmentation only where needed (from coins).
            final = torch.where(mask, result2, result)
            result = final.clone()
        return result
