"""
"XFeat: Accelerated Features for Lightweight Image Matching, CVPR 2024."
https://www.verlab.dcc.ufmg.br/descriptors/xfeat_cvpr24/

This script implements color + geometric transformations using Kornia.
Given a dataset of random real unlabeled images, we apply photometric transformations,
homography warps and also TPS warps. It also handles black borders by
pasting a random background image.
"""

import torch
from torch import nn
from torch.utils.data import Dataset
import torch.utils.data as data
from torchvision import transforms
import torch.nn.functional as F

import cv2
import kornia
import kornia.augmentation as K
from kornia.geometry.transform import get_tps_transform as findTPS
from kornia.geometry.transform import warp_points_tps, warp_image_tps

import glob
import random
import tqdm

import numpy as np
import pdb
import time

random.seed(0)
torch.manual_seed(0)


def generateRandomTPS(shape, grid=(8, 6), GLOBAL_MULTIPLIER=0.3, prob=0.5):

    h, w = shape
    sh, sw = h / grid[0], w / grid[1]
    src = torch.dstack(torch.meshgrid(torch.arange(0, h + sh, sh), torch.arange(0, w + sw, sw), indexing="ij"))

    offsets = torch.rand(grid[0] + 1, grid[1] + 1, 2) - 0.5
    offsets *= torch.tensor([sh / 2, sw / 2]).view(1, 1, 2) * min(0.97, 2.0 * GLOBAL_MULTIPLIER)
    dst = src + offsets if np.random.uniform() < prob else src

    src, dst = src.view(1, -1, 2), dst.view(1, -1, 2)
    src = (src / torch.tensor([h, w]).view(1, 1, 2)) * 2 - 1.0
    dst = (dst / torch.tensor([h, w]).view(1, 1, 2)) * 2 - 1.0
    weights, A = findTPS(dst, src)

    return src, weights, A


def generateRandomHomography(shape, GLOBAL_MULTIPLIER=0.3):
    # Generate random in-plane rotation [-theta,+theta]
    theta = np.radians(np.random.uniform(-30, 30))

    # Generate random scale in both x and y
    scale_x, scale_y = np.random.uniform(0.35, 1.2, 2)

    # Generate random translation shift
    tx, ty = -shape[1] / 2.0, -shape[0] / 2.0
    txn, tyn = np.random.normal(0, 120.0 * GLOBAL_MULTIPLIER, 2)

    c, s = np.cos(theta), np.sin(theta)

    # Affine coeffs
    sx, sy = np.random.normal(0, 0.6 * GLOBAL_MULTIPLIER, 2)

    # Projective coeffs
    p1, p2 = np.random.normal(0, 0.006 * GLOBAL_MULTIPLIER, 2)

    # Build Homography from parmeterizations
    H_t = np.array(((1, 0, tx), (0, 1, ty), (0, 0, 1)))  # t

    H_r = np.array(((c, -s, 0), (s, c, 0), (0, 0, 1)))  # rotation,
    H_a = np.array(((1, sy, 0), (sx, 1, 0), (0, 0, 1)))  # affine
    H_p = np.array(((1, 0, 0), (0, 1, 0), (p1, p2, 1)))  # projective

    H_s = np.array(((scale_x, 0, 0), (0, scale_y, 0), (0, 0, 1)))  # scale
    H_b = np.array(((1.0, 0, -tx + txn), (0, 1, -ty + tyn), (0, 0, 1)))  # t_back,

    # H = H_e * H_s * H_a * H_p
    H = np.dot(np.dot(np.dot(np.dot(np.dot(H_b, H_s), H_p), H_a), H_r), H_t)

    return H


class AugmentationPipe(nn.Module):
    def __init__(
        self,
        device,
        load_dataset=True,
        img_dir="/homeLocal/guipotje/sshfs/datasets/coco_20k",
        warp_resolution=(1200, 900),
        out_resolution=(400, 300),
        sides_crop=0.2,
        max_num_imgs=50,
        num_test_imgs=10,
        batch_size=1,
        photometric=True,
        geometric=True,
        reload_step=1_000,
    ):
        super(AugmentationPipe, self).__init__()
        self.half = 16
        self.device = device

        self.dims = warp_resolution
        self.batch_size = batch_size
        self.out_resolution = out_resolution
        self.sides_crop = sides_crop
        self.max_num_imgs = max_num_imgs
        self.num_test_imgs = num_test_imgs
        self.dims_t = (
            torch.tensor(
                [
                    int(self.dims[0] * (1.0 - self.sides_crop)) - int(self.dims[0] * self.sides_crop) - 1,
                    int(self.dims[1] * (1.0 - self.sides_crop)) - int(self.dims[1] * self.sides_crop) - 1,
                ]
            )
            .float()
            .to(device)
            .view(1, 1, 2)
        )
        self.dims_s = (
            torch.tensor([self.dims_t[0, 0, 0] / out_resolution[0], self.dims_t[0, 0, 1] / out_resolution[1]])
            .float()
            .to(device)
            .view(1, 1, 2)
        )

        self.all_imgs = glob.glob(img_dir + "/*.jpg") + glob.glob(img_dir + "/*.png")

        self.photometric = photometric
        self.geometric = geometric
        self.cnt = 1
        self.reload_step = reload_step

        list_augmentation = [
            # kornia.augmentation.RandomChannelShuffle(p=0.5),
            kornia.augmentation.ColorJitter(0.15, 0.15, 0.15, 0.15, p=1.0),
            kornia.augmentation.RandomEqualize(p=0.4),
            kornia.augmentation.RandomGaussianBlur(p=0.3, sigma=(2.0, 2.0), kernel_size=(7, 7)),
        ]
        if photometric is False:
            list_augmentation = []

        self.aug_list = kornia.augmentation.ImageSequential(*list_augmentation)

        if len(self.all_imgs) < 10:
            raise RuntimeError("Couldnt find enough images to train. Please check the path: ", img_dir)

        if load_dataset:
            print("[Synthetic] Found a total of ", len(self.all_imgs), " images for training..")

            if len(self.all_imgs) - num_test_imgs < max_num_imgs:
                raise RuntimeError("Error: test set overlaps with training set! Decrease number of test imgs")

            self.load_imgs()

            self.TPS = True

    def load_imgs(self):

        random.shuffle(self.all_imgs)

        train = []
        fast = cv2.FastFeatureDetector_create(30)
        for p in tqdm.tqdm(self.all_imgs[: self.max_num_imgs], desc="loading train"):
            im = cv2.imread(p)
            halfH, halfW = im.shape[0] // 2, im.shape[1] // 2
            if halfH > halfW:
                im = np.rot90(im)
                halfH, halfW = halfW, halfH

            if im.shape[0] != self.dims[1] or im.shape[1] != self.dims[0]:
                im = cv2.resize(im, self.dims)

            train.append(np.copy(im))

        self.train = train

        self.test = [
            cv2.resize(cv2.imread(p), self.dims)
            for p in tqdm.tqdm(self.all_imgs[-self.num_test_imgs :], desc="loading test")
        ]

    def norm_pts_grid(self, x):
        if len(x.size()) == 2:
            return (x.view(1, -1, 2) * self.dims_s / self.dims_t) * 2.0 - 1
        return (x * self.dims_s / self.dims_t) * 2.0 - 1

    def denorm_pts_grid(self, x):
        if len(x.size()) == 2:
            return ((x.view(1, -1, 2) + 1) / 2.0) / self.dims_s * self.dims_t
        return ((x + 1) / 2.0) / self.dims_s * self.dims_t

    def rnd_kps(self, shape, n=256):
        h, w = shape
        kps = torch.rand(size=(3, n)).to(self.device)
        kps[0, :] *= w
        kps[1, :] *= h
        kps[2, :] = 1.0

        return kps

    def warp_points(self, H, pts):
        scale = self.dims_s.view(-1, 2)
        offset = torch.tensor(
            [int(self.dims[0] * self.sides_crop), int(self.dims[1] * self.sides_crop)], device=pts.device
        ).float()
        pts = pts * scale + offset
        pts = torch.vstack([pts.t(), torch.ones(1, pts.shape[0], device=pts.device)])
        warped = torch.matmul(H, pts)
        warped = warped / warped[2, ...]
        warped = warped.t()[:, :2]
        return (warped - offset) / scale

    @torch.inference_mode()
    def forward(self, x, difficulty=0.3, TPS=False, prob_deformation=0.5, test=False):
        """
        Perform augmentation to a batch of images.

        input:
            x -> torch.Tensor(B, C, H, W): rgb images
            difficulty -> float: level of difficulty, 0.1 is medium, 0.3 is already pretty hard
            tps -> bool: Wether to apply non-rigid deformations in images
            prob_deformation -> float: probability to apply a deformation

        return:
            'output'    ->   torch.Tensor(B, C, H, W): rgb images
            Tuple:
                'H'       ->   torch.Tensor(3,3): homography matrix
                'mask'  ->     torch.Tensor(B, H, W): mask of valid pixels after warp
                (deformation only)
                src, weights, A are parameters from a TPS warp (all torch.Tensors)

        """

        if self.cnt % self.reload_step == 0:
            self.load_imgs()

        if self.geometric is False:
            difficulty = 0.0

        with torch.no_grad():
            x = (x / 255.0).to(self.device)
            b, c, h, w = x.shape
            shape = (h, w)

            ######## Geometric Transformations

            H = torch.tensor(
                np.array([generateRandomHomography(shape, difficulty) for b in range(self.batch_size)]),
                dtype=torch.float32,
            ).to(self.device)

            output = kornia.geometry.transform.warp_perspective(x, H, dsize=shape, padding_mode="zeros")

            # crop % of image boundaries each side to reduce invalid pixels after warps
            low_h = int(h * self.sides_crop)
            low_w = int(w * self.sides_crop)
            high_h = int(h * (1.0 - self.sides_crop))
            high_w = int(w * (1.0 - self.sides_crop))
            output = output[..., low_h:high_h, low_w:high_w]
            x = x[..., low_h:high_h, low_w:high_w]

            # apply TPS if desired:
            if TPS:
                src, weights, A = None, None, None
                for b in range(self.batch_size):
                    b_src, b_weights, b_A = generateRandomTPS(shape, (8, 6), difficulty, prob=prob_deformation)
                    b_src, b_weights, b_A = b_src.to(self.device), b_weights.to(self.device), b_A.to(self.device)

                    if src is None:
                        src, weights, A = b_src, b_weights, b_A
                    else:
                        src = torch.cat((b_src, src))
                        weights = torch.cat((b_weights, weights))
                        A = torch.cat((b_A, A))

                output = warp_image_tps(output, src, weights, A)

            output = F.interpolate(output, self.out_resolution[::-1], mode="nearest")
            x = F.interpolate(x, self.out_resolution[::-1], mode="nearest")

            mask = ~torch.all(output == 0, dim=1, keepdim=True)
            mask = mask.expand(-1, 3, -1, -1)

            # Make-up invalid regions with texture from the batch
            rv = 1 if not TPS else 2
            output_shifted = torch.roll(x, rv, 0)
            output[~mask] = output_shifted[~mask]
            mask = mask[:, 0, :, :]

            ######## Photometric Transformations
            output = self.aug_list(output)

            b, c, h, w = output.shape
            # Correlated Gaussian Noise
            if np.random.uniform() > 0.5 and self.photometric:
                noise = F.interpolate(torch.randn_like(output) * (10 / 255), (h // 2, w // 2))
                noise = F.interpolate(noise, (h, w), mode="bicubic")
                output = torch.clip(output + noise, 0.0, 1.0)

            # Random shadows
            if np.random.uniform() > 0.6 and self.photometric:
                noise = torch.rand((b, 1, h // 64, w // 64), device=self.device) * 1.3
                noise = torch.clip(noise, 0.25, 1.0)
                noise = F.interpolate(noise, (h, w), mode="bicubic")
                noise = noise.expand(-1, 3, -1, -1)
                output *= noise
                output = torch.clip(output, 0.0, 1.0)

            self.cnt += 1

        if TPS:
            return output, (H, src, weights, A, mask)
        else:
            return output, (H, mask)

    def get_correspondences(self, kps_target, T):
        H, H2, src, W, A = T
        undeformed = self.denorm_pts_grid(warp_points_tps(self.norm_pts_grid(kps_target), src, W, A)).view(-1, 2)

        warped_to_src = self.warp_points(H @ torch.inverse(H2), undeformed)

        return warped_to_src
