import cv2
import numpy as np
import torch
import os

import sys

ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, ROOT_DIR)

from superpoint import SuperPoint


def resize(img, resize):
    img_h, img_w = img.shape[0], img.shape[1]
    cur_size = max(img_h, img_w)
    if len(resize) == 1:
        scale1, scale2 = resize[0] / cur_size, resize[0] / cur_size
    else:
        scale1, scale2 = resize[0] / img_h, resize[1] / img_w
    new_h, new_w = int(img_h * scale1), int(img_w * scale2)
    new_img = cv2.resize(img.astype("float32"), (new_w, new_h)).astype("uint8")
    scale = np.asarray([scale2, scale1])
    return new_img, scale


class ExtractSIFT:
    def __init__(self, config, root=True):
        self.num_kp = config["num_kpt"]
        self.contrastThreshold = config["det_th"]
        self.resize = config["resize"]
        self.root = root

    def run(self, img_path):
        self.sift = cv2.xfeatures2d.SIFT_create(
            nfeatures=self.num_kp, contrastThreshold=self.contrastThreshold
        )
        img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
        scale = [1, 1]
        if self.resize[0] != -1:
            img, scale = resize(img, self.resize)
        cv_kp, desc = self.sift.detectAndCompute(img, None)
        kp = np.array(
            [
                [_kp.pt[0] / scale[1], _kp.pt[1] / scale[0], _kp.response]
                for _kp in cv_kp
            ]
        )  # N*3
        index = np.flip(np.argsort(kp[:, 2]))
        kp, desc = kp[index], desc[index]
        if self.root:
            desc = np.sqrt(
                abs(desc / (np.linalg.norm(desc, axis=-1, ord=1)[:, np.newaxis] + 1e-8))
            )
        return kp[: self.num_kp], desc[: self.num_kp]


class ExtractSuperpoint(object):
    def __init__(self, config):
        default_config = {
            "descriptor_dim": 256,
            "nms_radius": 4,
            "detection_threshold": config["det_th"],
            "max_keypoints": config["num_kpt"],
            "remove_borders": 4,
            "model_path": "../weights/sp/superpoint_v1.pth",
        }
        self.superpoint_extractor = SuperPoint(default_config)
        self.superpoint_extractor.eval(), self.superpoint_extractor.cuda()
        self.num_kp = config["num_kpt"]
        if "padding" in config.keys():
            self.padding = config["padding"]
        else:
            self.padding = False
        self.resize = config["resize"]

    def run(self, img_path):
        img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
        scale = 1
        if self.resize[0] != -1:
            img, scale = resize(img, self.resize)
        with torch.no_grad():
            result = self.superpoint_extractor(
                torch.from_numpy(img / 255.0).float()[None, None].cuda()
            )
        score, kpt, desc = (
            result["scores"][0],
            result["keypoints"][0],
            result["descriptors"][0],
        )
        score, kpt, desc = score.cpu().numpy(), kpt.cpu().numpy(), desc.cpu().numpy().T
        kpt = np.concatenate([kpt / scale, score[:, np.newaxis]], axis=-1)
        # padding randomly
        if self.padding:
            if len(kpt) < self.num_kp:
                res = int(self.num_kp - len(kpt))
                pad_x, pad_desc = np.random.uniform(size=[res, 2]) * (
                    img.shape[0] + img.shape[1]
                ) / 2, np.random.uniform(size=[res, 256])
                pad_kpt, pad_desc = (
                    np.concatenate([pad_x, np.zeros([res, 1])], axis=-1),
                    pad_desc / np.linalg.norm(pad_desc, axis=-1)[:, np.newaxis],
                )
                kpt, desc = np.concatenate([kpt, pad_kpt], axis=0), np.concatenate(
                    [desc, pad_desc], axis=0
                )
        return kpt, desc