import importlib
import math
import os
from pathlib import Path
from typing import Optional, Tuple

import cv2
import numpy as np
import requests
import torch

import kornia as K


def read_img_from_url(url: str, resize_to: Optional[Tuple[int, int]] = None) -> torch.Tensor:
    # perform request
    response = requests.get(url).content
    # convert to array of ints
    nparr = np.frombuffer(response, np.uint8)
    # convert to image array and resize
    img: np.ndarray = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)[..., :3]
    # convert the image to a tensor
    img_t: torch.Tensor = K.utils.image_to_tensor(img, keepdim=False)  # 1xCxHXW
    img_t = img_t.float() / 255.0
    if resize_to is None:
        img_t = K.geometry.resize(img_t, 184)
    else:
        img_t = K.geometry.resize(img_t, resize_to)
    return img_t


def main():
    # load the images
    BASE_IMAGE_URL1: str = "https://raw.githubusercontent.com/kornia/data/main/panda.jpg"  # augmentation
    BASE_IMAGE_URL2: str = "https://raw.githubusercontent.com/kornia/data/main/simba.png"  # color
    BASE_IMAGE_URL3: str = "https://raw.githubusercontent.com/kornia/data/main/girona.png"  # enhance
    BASE_IMAGE_URL4: str = "https://raw.githubusercontent.com/kornia/data/main/baby_giraffe.png"  # morphology
    BASE_IMAGE_URL5: str = "https://raw.githubusercontent.com/kornia/data/main/persistencia_memoria.jpg"  # filters
    BASE_IMAGE_URL6: str = "https://raw.githubusercontent.com/kornia/data/main/delorean.png"  # geometry
    OUTPUT_PATH = Path(__file__).absolute().parent / "source/_static/img"

    os.makedirs(OUTPUT_PATH, exist_ok=True)
    print(f"Pointing images to path {OUTPUT_PATH}.")
    img1 = read_img_from_url(BASE_IMAGE_URL1)
    img2 = read_img_from_url(BASE_IMAGE_URL2, img1.shape[-2:])
    img3 = read_img_from_url(BASE_IMAGE_URL3, img1.shape[-2:])
    img4 = read_img_from_url(BASE_IMAGE_URL4)
    img5 = read_img_from_url(BASE_IMAGE_URL5, (234, 320))
    img6 = read_img_from_url(BASE_IMAGE_URL6)

    # TODO: make this more generic for modules out of kornia.augmentation
    # Dictionary containing the transforms to generate the sample images:
    # Key: Name of the transform class.
    # Value: (parameters, num_samples, seed)
    mod = importlib.import_module("kornia.augmentation")
    augmentations_list: dict = {
        "CenterCrop": ((184, 184), 1, 2018),
        "ColorJitter": ((0.3, 0.3, 0.3, 0.3), 2, 2018),
        "RandomAffine": (((-15.0, 20.0), (0.1, 0.1), (0.7, 1.3), 20), 2, 2019),
        "RandomBoxBlur": (((7, 7),), 1, 2020),
        "RandomCrop": ((img1.shape[-2:], (50, 50)), 2, 2020),
        "RandomChannelShuffle": ((), 1, 2020),
        "RandomElasticTransform": (((63, 63), (32, 32), (2.0, 2.0)), 2, 2018),
        "RandomEqualize": ((), 1, 2020),
        "RandomErasing": (((0.2, 0.4), (0.3, 1 / 0.3)), 2, 2017),
        "RandomFisheye": ((torch.tensor([-0.3, 0.3]), torch.tensor([-0.3, 0.3]), torch.tensor([0.9, 1.0])), 2, 2020),
        "RandomGaussianBlur": (((3, 3), (0.1, 2.0)), 1, 2020),
        "RandomGaussianNoise": ((0.0, 0.05), 1, 2020),
        "RandomGrayscale": ((), 1, 2020),
        "RandomHorizontalFlip": ((), 1, 2020),
        "RandomInvert": ((), 1, 2020),
        "RandomMotionBlur": ((7, 35.0, 0.5), 2, 2020),
        "RandomPerspective": ((0.2,), 2, 2020),
        "RandomPlanckianJitter": ((), 2, 2022),
        "RandomPosterize": (((1, 4),), 2, 2016),
        "RandomResizedCrop": ((img1.shape[-2:], (1.0, 2.0), (1.0, 2.0)), 2, 2020),
        "RandomRotation": ((45.0,), 2, 2019),
        "RandomSharpness": ((16.0,), 1, 2019),
        "RandomSolarize": ((0.2, 0.2), 2, 2019),
        "RandomVerticalFlip": ((), 1, 2020),
        "RandomThinPlateSpline": ((), 1, 2020),
    }

    # ITERATE OVER THE TRANSFORMS
    for aug_name, (args, num_samples, seed) in augmentations_list.items():
        img_in = img1.repeat(num_samples, 1, 1, 1)
        # dynamically create the class instance
        cls = getattr(mod, aug_name)
        aug = cls(*args, p=1.0)
        # set seed
        torch.manual_seed(seed)
        # apply the augmentaiton to the image and concat
        out = aug(img_in)

        if aug_name == "CenterCrop":
            h, w = img1.shape[-2:]
            h_new, w_new = out.shape[-2:]
            h_dif, w_dif = int(h - h_new), int(w - w_new)
            out = torch.nn.functional.pad(out, (w_dif // 2, w_dif // 2, 0, h_dif))

        out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))], dim=-1)
        # save the output image
        out_np = K.utils.tensor_to_image((out * 255.0).byte())
        cv2.imwrite(str(OUTPUT_PATH / f"{aug_name}.png"), out_np)
        sig = f"{aug_name}({', '.join([str(a) for a in args])}, p=1.0)"
        print(f"Generated image example for {aug_name}. {sig}")

    mod = importlib.import_module("kornia.augmentation")
    mix_augmentations_list: dict = {
        "RandomMixUp": (((0.3, 0.4),), 2, 20),
        "RandomCutMix": ((img1.shape[-2], img1.shape[-1]), 2, 2019),
    }
    # ITERATE OVER THE TRANSFORMS
    for aug_name, (args, num_samples, seed) in mix_augmentations_list.items():
        img_in = torch.cat([img1, img2])
        # dynamically create the class instance
        cls = getattr(mod, aug_name)
        aug = cls(*args, p=1.0)
        # set seed
        torch.manual_seed(seed)
        # apply the augmentaiton to the image and concat
        out, _ = aug(img_in, torch.tensor([0, 1]))
        out = torch.cat([img_in[0], img_in[1], *(out[i] for i in range(out.size(0)))], dim=-1)
        # save the output image
        out_np = K.utils.tensor_to_image((out * 255.0).byte())
        cv2.imwrite(str(OUTPUT_PATH / f"{aug_name}.png"), out_np)
        sig = f"{aug_name}({', '.join([str(a) for a in args])}, p=1.0)"
        print(f"Generated image example for {aug_name}. {sig}")

    mod = importlib.import_module("kornia.color")
    color_transforms_list: dict = {
        "grayscale_to_rgb": ((), 3),
        "rgb_to_bgr": ((), 1),
        "rgb_to_grayscale": ((), 1),
        "rgb_to_hsv": ((), 1),
        "rgb_to_hls": ((), 1),
        "rgb_to_luv": ((), 1),
        "rgb_to_lab": ((), 1),
        # "rgb_to_rgba": ((1.,), 1),
        "rgb_to_xyz": ((), 1),
        "rgb_to_ycbcr": ((), 1),
        "rgb_to_yuv": ((), 1),
        "rgb_to_linear_rgb": ((), 1),
    }
    # ITERATE OVER THE TRANSFORMS
    for fn_name, (args, num_samples) in color_transforms_list.items():
        # import function and apply
        fn = getattr(mod, fn_name)
        if fn_name == "grayscale_to_rgb":
            out = fn(K.color.rgb_to_grayscale(img2), *args)
        else:
            out = fn(img2, *args)
        # perform normalization to visualize
        if fn_name == "rgb_to_lab":
            out = out[:, :1] / 100.0
        elif fn_name == "rgb_to_hsv":
            out[:, :1] = out[:, :1] / 2 * math.pi
        elif fn_name == "rgb_to_luv":
            out = out[:, :1] / 116.0
        # repeat channels for grayscale
        if out.shape[1] != 3:
            out = out.repeat(1, 3, 1, 1)
        # save the output image
        if fn_name == "grayscale_to_rgb":
            out = torch.cat(
                [K.color.rgb_to_grayscale(img2[0]).repeat(3, 1, 1), *(out[i] for i in range(out.size(0)))], dim=-1
            )
        else:
            out = torch.cat([img2[0], *(out[i] for i in range(out.size(0)))], dim=-1)
        out_np = K.utils.tensor_to_image((out * 255.0).byte())
        cv2.imwrite(str(OUTPUT_PATH / f"{fn_name}.png"), out_np)
        sig = f"{fn_name}({', '.join([str(a) for a in args])})"
        print(f"Generated image example for {fn_name}. {sig}")

    # korna.enhance module
    mod = importlib.import_module("kornia.enhance")
    transforms: dict = {
        "adjust_brightness": ((torch.tensor([0.25, 0.5]),), 2),
        "adjust_contrast": ((torch.tensor([0.65, 0.5]),), 2),
        "adjust_gamma": ((torch.tensor([0.85, 0.75]), 2.0), 2),
        "adjust_hue": ((torch.tensor([-math.pi / 4, math.pi / 4]),), 2),
        "adjust_saturation": ((torch.tensor([1.0, 2.0]),), 2),
        "solarize": ((torch.tensor([0.8, 0.5]), torch.tensor([-0.25, 0.25])), 2),
        "posterize": ((torch.tensor([4, 2]),), 2),
        "sharpness": ((torch.tensor([1.0, 2.5]),), 2),
        "equalize": ((), 1),
        "invert": ((), 1),
        "equalize_clahe": ((), 1),
        "add_weighted": ((0.75, 0.25, 2.0), 1),
    }
    # ITERATE OVER THE TRANSFORMS
    for fn_name, (args, num_samples) in transforms.items():
        img_in = img3.repeat(num_samples, 1, 1, 1)
        if fn_name == "add_weighted":
            args_in = (img_in, args[0], img2, args[1], args[2])
        else:
            args_in = (img_in, *args)
        # import function and apply
        fn = getattr(mod, fn_name)
        out = fn(*args_in)
        # save the output image
        out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))], dim=-1)
        out_np = K.utils.tensor_to_image((out * 255.0).byte())
        cv2.imwrite(str(OUTPUT_PATH / f"{fn_name}.png"), out_np)
        sig = f"{fn_name}({', '.join([str(a) for a in args])})"
        print(f"Generated image example for {fn_name}. {sig}")

    # korna.morphology module
    mod = importlib.import_module("kornia.morphology")
    kernel = torch.tensor([[0, 1, 0], [1, 1, 1], [0, 1, 0]])
    transforms: dict = {
        "dilation": ((kernel,), 1),
        "erosion": ((kernel,), 1),
        "opening": ((kernel,), 1),
        "closing": ((kernel,), 1),
        "gradient": ((kernel,), 1),
        "top_hat": ((kernel,), 1),
        "bottom_hat": ((kernel,), 1),
    }
    # ITERATE OVER THE TRANSFORMS
    for fn_name, (args, num_samples) in transforms.items():
        img_in = img4.repeat(num_samples, 1, 1, 1)
        args_in = (img_in, *args)
        # import function and apply
        # import pdb;pdb.set_trace()
        fn = getattr(mod, fn_name)
        out = fn(*args_in)
        # save the output image
        out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))], dim=-1)
        out_np = K.utils.tensor_to_image((out * 255.0).byte())
        cv2.imwrite(str(OUTPUT_PATH / f"{fn_name}.png"), out_np)
        sig = f"{fn_name}({', '.join([str(a) for a in args])})"
        print(f"Generated image example for {fn_name}. {sig}")

    # korna.filters module
    mod = importlib.import_module("kornia.filters")
    kernel = torch.tensor([[0, 1, 0], [1, 1, 1], [0, 1, 0]])
    transforms: dict = {
        "box_blur": (((5, 5),), 1),
        "median_blur": (((5, 5),), 1),
        "gaussian_blur2d": (((5, 5), (1.5, 1.5)), 1),
        "motion_blur": ((5, 90.0, 1.0), 1),
        "max_blur_pool2d": ((5,), 1),
        "blur_pool2d": ((5,), 1),
        "unsharp_mask": (((5, 5), (1.5, 1.5)), 1),
        "laplacian": ((5,), 1),
        "sobel": ((), 1),
        "spatial_gradient": ((), 1),
        "canny": ((), 1),
    }
    # ITERATE OVER THE TRANSFORMS
    for fn_name, (args, num_samples) in transforms.items():
        img_in = img5.repeat(num_samples, 1, 1, 1)
        args_in = (img_in, *args)
        # import function and apply
        fn = getattr(mod, fn_name)
        out = fn(*args_in)
        if fn_name in ("max_blur_pool2d", "blur_pool2d"):
            out = K.geometry.resize(out, img_in.shape[-2:])
        if fn_name == "canny":
            out = out[1].repeat(1, 3, 1, 1)
        if isinstance(out, torch.Tensor):
            out = out.clamp(min=0.0, max=1.0)
        if fn_name in ("laplacian", "sobel", "spatial_gradient", "canny"):
            out = K.enhance.normalize_min_max(out)
        if fn_name == "spatial_gradient":
            out = out.permute(2, 1, 0, 3, 4).squeeze()
        # save the output image
        out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))], dim=-1)
        out_np = K.utils.tensor_to_image((out * 255.0).byte())
        cv2.imwrite(str(OUTPUT_PATH / f"{fn_name}.png"), out_np)
        sig = f"{fn_name}({', '.join([str(a) for a in args])})"
        print(f"Generated image example for {fn_name}. {sig}")

    # korna.geometry.transform module
    mod = importlib.import_module("kornia.geometry.transform")
    h, w = img6.shape[-2:]

    def _get_tps_args():
        src = torch.tensor([[[-1.0, -1.0], [-1.0, 1.0], [1.0, -1.0], [1.0, -1.0], [0.0, 0.0]]]).repeat(2, 1, 1)  # Bx5x2
        dst = src + torch.distributions.Uniform(-0.2, 0.2).rsample((2, 5, 2))
        kernel, affine = K.geometry.transform.get_tps_transform(dst, src)
        return src, kernel, affine

    transforms: dict = {
        "warp_affine": (
            (
                K.geometry.transform.get_affine_matrix2d(
                    translations=torch.zeros(2, 2),
                    center=(torch.tensor([w, h]) / 2).repeat(2, 1),
                    scale=torch.distributions.Uniform(0.5, 1.5).rsample((2, 2)),
                    angle=torch.tensor([-25.0, 25.0]),
                )[:, :2, :3],
                (h, w),
            ),
            2,
        ),
        "remap": (
            (
                *(K.utils.create_meshgrid(h, w, normalized_coordinates=True) - 0.25).unbind(-1),
                'bilinear',
                'zeros',
                True,
                True,
            ),
            1,
        ),
        "warp_image_tps": ((_get_tps_args()), 2),
        "rotate": ((torch.tensor([-15.0, 25.0]),), 2),
        "translate": ((torch.tensor([[10.0, -15], [50.0, -25.0]]),), 2),
        "scale": ((torch.tensor([[0.5, 1.25], [1.0, 1.5]]),), 2),
        "shear": ((torch.tensor([[0.1, -0.2], [-0.2, 0.1]]),), 2),
        "rot180": ((), 1),
        "hflip": ((), 1),
        "vflip": ((), 1),
        "resize": (((120, 220),), 1),
        "rescale": ((0.5,), 1),
        "elastic_transform2d": ((torch.rand(1, 2, h, w) * 2 - 1, (63, 63), (32, 32), (4.0, 4.0)), 1),
        "pyrdown": ((), 1),
        "pyrup": ((), 1),
        "build_pyramid": ((3,), 1),
    }
    # ITERATE OVER THE TRANSFORMS
    for fn_name, (args, num_samples) in transforms.items():
        img_in = img6.repeat(num_samples, 1, 1, 1)
        args_in = (img_in, *args)
        # import function and apply
        fn = getattr(mod, fn_name)
        out = fn(*args_in)
        if fn_name in ("resize", "rescale", "pyrdown", "pyrup"):
            h_new, w_new = out.shape[-2:]
            out = torch.nn.functional.pad(out, (0, (w - w_new), 0, (h - h_new)))
        if fn_name == "build_pyramid":
            _out = []
            for pyr in out[1:]:
                h_new, w_new = pyr.shape[-2:]
                out_tmp = torch.nn.functional.pad(pyr, (0, (w - w_new), 0, (h - h_new)))
                _out.append(out_tmp)
            out = torch.cat(_out)
        # save the output image
        out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))], dim=-1)
        out_np = K.utils.tensor_to_image((out * 255.0).byte())
        cv2.imwrite(str(OUTPUT_PATH / f"{fn_name}.png"), out_np)
        sig = f"{fn_name}({', '.join([str(a) for a in args])})"
        print(f"Generated image example for {fn_name}. {sig}")


if __name__ == "__main__":
    main()

[(19, 'numpy.frombuffer', 'np.frombuffer', True, 'import numpy as np\n'), (23, 'kornia.utils.image_to_tensor', 'K.utils.image_to_tensor', True, 'import kornia as K\n'), (42, 'os.makedirs', 'os.makedirs', False, 'import os\n'), (55, 'importlib.import_module', 'importlib.import_module', False, 'import importlib\n'), (108, 'importlib.import_module', 'importlib.import_module', False, 'import importlib\n'), (130, 'importlib.import_module', 'importlib.import_module', False, 'import importlib\n'), (176, 'importlib.import_module', 'importlib.import_module', False, 'import importlib\n'), (209, 'importlib.import_module', 'importlib.import_module', False, 'import importlib\n'), (210, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (236, 'importlib.import_module', 'importlib.import_module', False, 'import importlib\n'), (237, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (276, 'importlib.import_module', 'importlib.import_module', False, 'import importlib\n'), (17, 'requests.get', 'requests.get', False, 'import requests\n'), (21, 'cv2.imdecode', 'cv2.imdecode', False, 'import cv2\n'), (26, 'kornia.geometry.resize', 'K.geometry.resize', True, 'import kornia as K\n'), (28, 'kornia.geometry.resize', 'K.geometry.resize', True, 'import kornia as K\n'), (91, 'torch.manual_seed', 'torch.manual_seed', False, 'import torch\n'), (115, 'torch.cat', 'torch.cat', False, 'import torch\n'), (120, 'torch.manual_seed', 'torch.manual_seed', False, 'import torch\n'), (282, 'kornia.geometry.transform.get_tps_transform', 'K.geometry.transform.get_tps_transform', True, 'import kornia as K\n'), (99, 'torch.nn.functional.pad', 'torch.nn.functional.pad', False, 'import torch\n'), (122, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (259, 'kornia.geometry.resize', 'K.geometry.resize', True, 'import kornia as K\n'), (265, 'kornia.enhance.normalize_min_max', 'K.enhance.normalize_min_max', True, 'import kornia as K\n'), (332, 'torch.nn.functional.pad', 'torch.nn.functional.pad', False, 'import torch\n'), (339, 'torch.cat', 'torch.cat', False, 'import torch\n'), (66, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (66, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (66, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (150, 'kornia.color.rgb_to_grayscale', 'K.color.rgb_to_grayscale', True, 'import kornia as K\n'), (178, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (179, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (180, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (181, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (182, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (183, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (183, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (184, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (185, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (280, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (309, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (310, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (311, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (312, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (337, 'torch.nn.functional.pad', 'torch.nn.functional.pad', False, 'import torch\n'), (40, 'pathlib.Path', 'Path', False, 'from pathlib import Path\n'), (281, 'torch.distributions.Uniform', 'torch.distributions.Uniform', False, 'import torch\n'), (318, 'torch.rand', 'torch.rand', False, 'import torch\n'), (166, 'kornia.color.rgb_to_grayscale', 'K.color.rgb_to_grayscale', True, 'import kornia as K\n'), (289, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (292, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (300, 'kornia.utils.create_meshgrid', 'K.utils.create_meshgrid', True, 'import kornia as K\n'), (291, 'torch.distributions.Uniform', 'torch.distributions.Uniform', False, 'import torch\n'), (290, 'torch.tensor', 'torch.tensor', False, 'import torch\n')]