# LICENSE HEADER MANAGED BY add-license-header
#
# Copyright 2018 Kornia Team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from __future__ import annotations

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

import cv2
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import requests
import torch
import torch.nn.functional as F
from kornia_moons.feature import visualize_LAF

import kornia as K

mpl.use("Agg")


def download_tutorials_examples(download_infos: dict[str, str], directory: Path):
    URL_BASE = "https://raw.githubusercontent.com/kornia/tutorials/master/"
    for filename, path in download_infos.items():
        url = URL_BASE + path
        # perform request
        response = requests.get(url, timeout=60).content

        path = directory / filename
        with open(path, "wb") as fp:
            fp.write(response)


def read_img_from_url(url: str, resize_to: Optional[tuple[int, int]] = None, **resize_kwargs) -> torch.Tensor:
    # perform request
    response = requests.get(url, timeout=60).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.image.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, **resize_kwargs)
    else:
        img_t = K.geometry.resize(img_t, resize_to, **resize_kwargs)
    return img_t


def transparent_pad(src: torch.Tensor, shape: tuple[int, int]) -> torch.Tensor:
    """Apply a transparent pad to src (centerized) to match with shape (h, w)"""
    w_pad = abs(int(src.shape[-1] - shape[-1]) // 2)
    h_pad = abs(int(src.shape[-2] - shape[-2]) // 2)
    return F.pad(K.color.rgb_to_rgba(src, 1.0), (w_pad, w_pad, h_pad, h_pad), "constant", 0.0)


def draw_bbox_kpts(imgs: torch.Tensor, bboxes: torch.Tensor, keypoints: torch.Tensor) -> torch.Tensor:
    rectangle = torch.zeros(imgs.shape[0], imgs.shape[1], 4)
    rectangle[..., 0] = bboxes[..., 0]  # x1
    rectangle[..., 1] = bboxes[..., 1]  # y1
    rectangle[..., 2] = bboxes[..., 0] + bboxes[..., -2]  # x2
    rectangle[..., 3] = bboxes[..., 1] + bboxes[..., -1]  # y2
    color = torch.tensor([1, 0, 0]).repeat(imgs.shape[0], imgs.shape[1], 1)
    imgs_draw = K.image.draw_rectangle(imgs, rectangle, color=color)

    rectangle2 = torch.zeros(imgs.shape[0], imgs.shape[1], 4)
    for n in range(keypoints.shape[-2]):
        rectangle2[..., n, 0] = keypoints[..., n, 0] - 2
        rectangle2[..., n, 1] = keypoints[..., n, 1] - 2
        rectangle2[..., n, 2] = keypoints[..., n, 0] + 2
        rectangle2[..., n, 3] = keypoints[..., n, 1] + 2
    color = torch.tensor([0, 0, 1]).repeat(imgs.shape[0], imgs.shape[1], 1)
    imgs_draw = K.utils.draw_rectangle(imgs_draw, rectangle2, color=color, fill=True)

    return imgs_draw


def main():
    # Download the tutorial examples for the main docs
    # Note: Training API examples (image_classifier, object_detection, semantic_segmentation) removed
    # as they depend on kornia.x which has been removed
    URLS_TUTORIALS_EXAMPLES = {}

    OUTPUT_PATH_SCRIPTS = Path(__file__).absolute().parent / "source/_static/scripts/"

    os.makedirs(OUTPUT_PATH_SCRIPTS, exist_ok=True)
    print(f"Downloading script examples from kornia/tutorials. Saving into the path {OUTPUT_PATH_SCRIPTS}.")
    download_tutorials_examples(URLS_TUTORIALS_EXAMPLES, OUTPUT_PATH_SCRIPTS)

    # 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
    hash1 = "8b98f44abbe92b7a84631ed06613b08fee7dae14"
    BASE_IMAGEOUTDOOR_URL7: str = f"https://github.com/kornia/data_test/raw/{hash1}/knchurch_disk.pt"  # image matching
    BASE_IMAGEOUTDOOR_URL8: str = (  # Response functions
        "https://github.com/kornia/data/raw/main/kornia_banner_pixie.png"
    )
    MASK_IMAGE_URL2: str = "https://raw.githubusercontent.com/kornia/data/main/simba_mask.png"

    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)
    img_kornia = read_img_from_url(BASE_IMAGEOUTDOOR_URL8)

    # Read the masks as (B, H, W)
    mask2 = read_img_from_url(MASK_IMAGE_URL2, img1.shape[-2:], interpolation="nearest")
    mask2 = mask2.median(dim=1)[0]

    # 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),
        "ColorJiggle": ((0.3, 0.3, 0.3, 0.3), 2, 2018),
        "ColorJitter": ((0.3, 0.3, 0.3, 0.3), 2, 2022),
        "PadTo": (((220, 450),), 1, 2022),
        "RandomAffine": (((-15.0, 20.0), (0.1, 0.1), (0.7, 1.3), 20), 2, 2019),
        "RandomBoxBlur": (((7, 7),), 1, 2020),
        "RandomBrightness": (((0.0, 1.0),), 2, 2022),
        "RandomContrast": (((0.0, 1.0),), 2, 2022),
        "RandomCrop": ((img1.shape[-2:], (50, 50)), 2, 2020),
        "RandomChannelDropout": ((), 1, 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),
        "RandomGamma": (((0.0, 1.0),), 2, 2022),
        "RandomGaussianBlur": (((3, 3), (0.1, 2.0)), 1, 2020),
        "RandomGaussianIllumination": (((0.5, 0.5), (0.5, 0.5), (0.5, 0.5), (-1.0, 1.0)), 2, 2021),
        "RandomGaussianNoise": ((0.0, 0.05), 1, 2020),
        "RandomGrayscale": ((), 1, 2020),
        "RandomHue": (((-0.5, 0.5),), 2, 2022),
        "RandomHorizontalFlip": ((), 1, 2020),
        "RandomInvert": ((), 1, 2020),
        "RandomJPEG": (((1.0, 5.0),), 1, 2024),
        "RandomLinearCornerIllumination": (((0.5, 0.5), (-1.0, 1.0)), 2, 2021),
        "RandomLinearIllumination": (((0.5, 0.5), (-1.0, 1.0)), 2, 2021),
        "RandomMedianBlur": (((3, 3),), 1, 2023),
        "RandomMotionBlur": ((7, 35.0, 0.5), 2, 2020),
        "RandomPerspective": ((0.2,), 2, 2020),
        "RandomPlanckianJitter": ((), 2, 2022),
        "RandomPlasmaShadow": (((0.2, 0.5),), 2, 2022),
        "RandomPlasmaBrightness": ((), 2, 2022),
        "RandomPlasmaContrast": ((), 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),
        "RandomSaltAndPepperNoise": (((0.05, 0.5), (0.1, 0.7)), 2, 2024),
        "RandomSaturation": (((0.5, 5.0),), 2, 2022),
        "RandomSharpness": ((16.0,), 1, 2019),
        "RandomSolarize": ((0.2, 0.2), 2, 2019),
        "RandomVerticalFlip": ((), 1, 2020),
        "RandomThinPlateSpline": ((), 1, 2020),
        "RandomJigsaw": ((), 2, 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)
        try:
            aug = cls(*args, p=1.0)
        except TypeError:
            aug = cls(*args)

        # set seed
        torch.manual_seed(seed)
        if aug_name == "RandomJigsaw":  # make sure the image is dividable
            img_in = K.geometry.resize(img_in, (1020, 500))
        elif aug_name == "RandomJPEG":
            img_in = img_in[..., :176, :]
        # apply the augmentation to the image and concat
        out = aug(img_in)

        # save ori image to concatenate into the out image
        ori = img_in[0]
        if aug_name == "CenterCrop":
            # Convert to RGBA, and center the output image with transparent pad
            out = transparent_pad(out, tuple(img1[-2:].shape))
            ori = K.color.rgb_to_rgba(ori, 1.0)  # To match the dims
        elif aug_name == "PadTo":
            # Convert to RGBA, and center the original image with transparent pad
            ori = transparent_pad(img_in[0], tuple(out.shape[-2:]))
            out = K.color.rgb_to_rgba(out, 1.0)  # To match the dims

        out = torch.cat([ori, *(out[i] for i in range(out.size(0)))], dim=-1)
        # save the output image
        out_np = K.image.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}")

    mix_augmentations_list = {"RandomMixUpV2": ((), 2, 20), "RandomCutMixV2": ((), 2, 2019), "PatchMix": ((), 2, 2024)}
    # ITERATE OVER THE TRANSFORMS
    for aug_name, (args, _, 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 augmentation to the image and concat
        # PatchMix returns (B, C, H, W); index [0] to get (C, H, W) for cat
        if aug_name == "PatchMix":
            img_aug = aug(img_in)[0]
        else:
            img_aug, _ = aug(img_in, torch.tensor([0, 1]))

        output = torch.cat([img_in[0], img_in[1], img_aug], dim=-1)
        # save the output image
        out_np = K.utils.tensor_to_image((output * 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}")

    mask_augmentations_list = {"RandomTransplantation": (([0],), 0)}
    # ITERATE OVER THE TRANSFORMS
    for aug_name, (args, seed) in mask_augmentations_list.items():
        img_in = torch.cat([img1, img2])
        mask_in = torch.cat([torch.zeros_like(mask2), mask2])

        # dynamically create the class instance
        cls = getattr(mod, aug_name)
        aug = cls(*args, p=1.0)
        # set seed
        torch.manual_seed(seed)
        # apply the augmentation to the image and concat
        img_aug, _ = aug(img_in, mask_in)

        output = torch.cat([img_in[0], img_in[1], img_aug[0]], dim=-1)
        # save the output image
        out_np = K.utils.tensor_to_image((output * 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}")

    # Containers
    aug_container_list = {
        "AugmentationSequential": (
            {
                "args": (
                    K.augmentation.ColorJitter(0.1, 0.1, 0.1, 0.1, p=1.0),
                    K.augmentation.RandomAffine(360, [0.1, 0.1], [0.7, 1.2], [30.0, 50.0], p=1.0),
                    K.augmentation.RandomPerspective(0.5, p=1.0),
                ),
                "data_keys": ["input", "bbox_xywh", "keypoints"],
            },
            (
                torch.tensor([[[125, 5, 115, 80]]], dtype=torch.float32),  # bbox
                torch.tensor([[[166, 42], [197, 42]]], dtype=torch.float32),  # keypoints
            ),
            2,
            2023,
        ),
        "PatchSequential": (
            {
                "args": (
                    K.augmentation.ColorJitter(0.2, 0.1, 0.1, 0.1, p=1),
                    K.augmentation.RandomAffine(10, [0.1, 0.2], [0.7, 1.2], [0.0, 15.0], p=1),
                    K.augmentation.RandomPerspective(0.3, p=1),
                    K.augmentation.RandomSolarize(0.01, 0.05, p=0.6),
                ),
                "grid_size": (2, 2),
                "same_on_batch": False,
                "patchwise_apply": False,
            },
            (),
            2,
            2023,
        ),
    }
    for aug_name, (args, labels, num_samples, seed) in aug_container_list.items():
        img_in = img1.repeat(num_samples, 1, 1, 1)
        cls = getattr(mod, aug_name)
        tfms = args.pop("args")
        augs = cls(*tfms, **args)

        # set seed
        torch.manual_seed(seed)
        if aug_name == "PatchSequential":
            out = augs(img_in)
            inp = img_in
        else:
            labels = (labels[0].expand(num_samples, -1, -1), labels[1].expand(num_samples, -1, -1))

            out = augs(img_in, *labels)
            out = draw_bbox_kpts(out[0], out[1].int(), out[2].int())
            inp = draw_bbox_kpts(img_in, labels[0].int(), labels[1].int())
        output = torch.cat([inp[0], *(out[i] for i in range(out.size(0)))], dim=-1)

        # save the output image
        out_np = K.utils.tensor_to_image((output * 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),
        "apply_colormap": ((K.color.ColorMap("autumn", 256),), 1),
        "ApplyColorMap": ((K.color.ColorMap("winter", 256),), 1),
    }
    # ITERATE OVER THE TRANSFORMS
    for fn_name, (args, _) 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)
        elif fn_name == "apply_colormap":
            gray_image = (K.color.rgb_to_grayscale(img2) * 255.0).round()
            out = K.color.rgb_to_bgr(fn(gray_image, *args))
        elif fn_name == "ApplyColorMap":
            gray_image = (K.color.rgb_to_grayscale(img2) * 255.0).round()
            out = K.color.rgb_to_bgr(fn(*args)(gray_image))
        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 in ("grayscale_to_rgb", "apply_colormap", "ApplyColorMap"):
            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.image.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.color.colormap
    colormaps_list = {"AUTUMN": (256,)}
    bar_img_gray = torch.arange(0, 256).repeat(1, 40, 1)  # 1x1x40x256
    bar_img = K.color.grayscale_to_rgb(bar_img_gray)
    # ITERATE OVER THE COLORMAPS
    for colormap_name, args in colormaps_list.items():
        cm = K.color.ColorMap(base=colormap_name, num_colors=args[0])
        out = K.color.rgb_to_bgr(K.color.apply_colormap(bar_img_gray, cm))[0]

        out = torch.cat([bar_img, out], dim=-1)

        out_np = K.image.tensor_to_image((out * 255.0).byte())
        cv2.imwrite(str(OUTPUT_PATH / f"{colormap_name}.png"), out_np)
        sig = f"{colormap_name}({', '.join([str(a) for a in args])})"
        print(f"Generated image example for {colormap_name}. {sig}")

    # Plot for all ColorMaps (ColorMapType.png)
    height_image = 40
    num_colors = 256
    num_columns = 3

    # 1 x height_image x num_colors
    input_tensor = torch.arange(start=0, end=num_colors, step=1).unsqueeze(0).repeat(1, height_image, 1)
    input_tensor = input_tensor.to("cpu").to(torch.float32)

    # Get colormap list
    colormap_list = K.color.ColorMapType.list()
    num_colormaps = len(colormap_list)
    # Calculate number of rows needed
    num_rows = (num_colormaps + num_columns - 1) // num_columns

    # Create figure and axis objects
    fig, axes = plt.subplots(num_rows, num_columns, figsize=(12, 8))
    for i, ax in enumerate(axes.flat):
        if i < num_colormaps:
            cmap = K.color.ColorMap(base=colormap_list[i], num_colors=num_colors)
            res = K.color.ApplyColorMap(colormap=cmap)(input_tensor)[0]
            ax.imshow(res.permute(1, 2, 0).numpy())
            ax.set_title(colormap_list[i], fontsize=12)
            ax.axis("off")
        else:
            fig.delaxes(ax)
    fig.tight_layout()
    fig.savefig(os.path.join(OUTPUT_PATH, "ColorMapType.png"), dpi=300)

    # 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),
        "jpeg_codec_differentiable": ((torch.tensor([50]),), 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 == "jpeg_codec_differentiable":
            img_in = img_in[..., :176, :]
        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.image.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.image.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 = {
        "bilateral_blur": (((11, 11), 0.1, (3, 3)), 1),
        "joint_bilateral_blur": (((11, 11), 0.1, (3, 3)), 1),
        "box_blur": (((5, 5),), 1),
        "median_blur": (((5, 5),), 1),
        "gaussian_blur2d": (((5, 5), (1.5, 1.5)), 1),
        "guided_blur": (((5, 5), 0.01), 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)
        if fn_name == "joint_bilateral_blur":
            guide = K.geometry.resize(img2.repeat(num_samples, 1, 1, 1), img_in.shape[-2:])
            args_in = (img_in, guide, *args)
        elif fn_name == "guided_blur":
            args_in = (img_in, img_in, *args)
        else:
            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()
        if fn_name == "joint_bilateral_blur":
            out = torch.cat([args_in[1], out], dim=-1)
        # save the output image
        out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))], dim=-1)
        out_np = K.image.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}")

    # kornia.filters.in_range
    mod = importlib.import_module("kornia.filters")
    transforms: dict = {
        "in_range": (((0.314, 0.2, 0.2), (0.47, 1.0, 1.0), True), 1),
    }
    # ITERATE OVER THE TRANSFORMS
    for fn_name, (args, _) in transforms.items():
        img_hsv = K.color.rgb_to_hsv(img1)
        h, s, v = torch.split(img_hsv, split_size_or_sections=1, dim=1)
        h = h / (2 * torch.pi)
        img_hsv = torch.cat((h, s, v), dim=1)
        args_in = (img_hsv, *args)
        fn = getattr(mod, fn_name)
        mask = fn(*args_in)
        filtered = img1 * mask
        mask = mask.repeat(1, img1.shape[1], 1, 1)
        # save the output image
        out = torch.cat([img1[0], mask[0], filtered[0]], dim=-1)
        out_np = K.image.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.geometry.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),
        "build_laplacian_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)

        if fn_name == "build_laplacian_pyramid":
            h_, w_ = out[0].shape[-2:]
            _out = [out[0]]
            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)))
                print(out_tmp.size())
                _out.append(out_tmp)
            out = torch.cat(_out)

        # save the output image
        if fn_name != "build_laplacian_pyramid":
            out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))], dim=-1)
        else:
            out = torch.cat([*(out[i] for i in range(out.size(0)))], dim=-1)
        out_np = K.image.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}")

    # Image Matching and local features
    img_matching_data = torch.hub.load_state_dict_from_url(BASE_IMAGEOUTDOOR_URL7, map_location=torch.device("cpu"))
    img_outdoor = img_matching_data["img2"]
    print("Generating local feature detections ")
    disk = K.feature.DISK.from_pretrained("depth")
    with torch.no_grad():
        disk_feat = disk(img_outdoor)[0]
        xy = disk_feat.keypoints.detach().cpu().numpy()
        cur_fname = str(OUTPUT_PATH / "disk_outdoor_depth.jpg")
        plt.figure()
        plt.imshow(K.tensor_to_image(img_outdoor))
        plt.scatter(xy[:, 0], xy[:, 1], 3, color="lime")
        plt.title('DISK("depth") keypoints')
        plt.savefig(cur_fname)
        plt.close()

    kah = K.feature.KeyNetAffNetHardNet(512).eval()
    with torch.no_grad():
        lafs, _resps, _descs = kah(K.color.rgb_to_grayscale(img_outdoor))
        fig1, ax = visualize_LAF(img_outdoor, lafs, color="lime", return_fig_ax=True, draw_ori=False)
        ax.set_title("KeyNetAffNet 512 LAFs")
        cur_fname = str(OUTPUT_PATH / "keynet_affnet.jpg")
        fig1.savefig(cur_fname)
        plt.close()

    keynet = K.feature.KeyNetDetector(True, 512).eval()
    with torch.no_grad():
        lafs, _resps = keynet(K.color.rgb_to_grayscale(img_outdoor))
        xy = K.feature.get_laf_center(lafs).detach().cpu().numpy().reshape(-1, 2)
        cur_fname = str(OUTPUT_PATH / "keynet.jpg")
        plt.figure()
        plt.imshow(K.tensor_to_image(img_outdoor))
        plt.scatter(xy[:, 0], xy[:, 1], 3, color="lime")
        plt.title("KeyNet 512 keypoints")
        plt.savefig(cur_fname)
        plt.close()

    # korna.feature module
    mod = importlib.import_module("kornia.feature")
    responses: list = [
        "harris_response",
        "gftt_response",
        "hessian_response",
        "dog_response_single",
        "KeyNet",
        "DISK",
        "ALIKED",
        "XFeat",
    ]
    # ITERATE OVER THE TRANSFORMS
    for fn_name in responses:
        # import function and apply
        img_in = K.color.rgb_to_grayscale(img_kornia)
        if fn_name == "KeyNet":
            fn = K.feature.KeyNet(True)
            out = fn(img_in)
        elif fn_name == "DISK":
            fn = K.feature.DISK.from_pretrained("depth")
            h, w = img_outdoor.shape[2:]
            pd_h = 32 - h % 32 if h % 32 > 0 else 0
            pd_w = 32 - w % 32 if w % 32 > 0 else 0
            img_in = torch.nn.functional.pad(img_outdoor, (0, pd_w, 0, pd_h), value=0.0)
            out, _ = fn.heatmap_and_dense_descriptors(img_in)
            out = K.color.grayscale_to_rgb(out)
            img_in = K.color.rgb_to_bgr(img_in)
        elif fn_name == "ALIKED":
            fn = K.feature.ALIKED.from_pretrained()
            with torch.no_grad():
                _, out = fn.extract_dense_map(img_outdoor)
            out = K.color.grayscale_to_rgb(out)
            img_in = K.color.rgb_to_bgr(img_outdoor)
        elif fn_name == "XFeat":
            fn = K.feature.XFeat.from_pretrained()
            with torch.no_grad():
                preprocessed, _, _ = fn._preprocess_tensor(img_outdoor)
                _, K1, _ = fn.net(preprocessed)
                out = K.feature.XFeat._get_kpts_heatmap(K1)
            # upsample heatmap from preprocessed size back to original image size
            out = torch.nn.functional.interpolate(out, img_outdoor.shape[-2:], mode="bilinear", align_corners=False)
            out = K.color.grayscale_to_rgb(out)
            img_in = K.color.rgb_to_bgr(img_outdoor)
        else:
            fn = getattr(mod, fn_name)
            out = fn(img_in)

        out = out - out.min()
        out = out / (1e-8 + out.max())

        # save the output image
        out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))], dim=-1)
        out_np = K.image.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 response function {fn_name}")


if __name__ == "__main__":
    main()