import cv2 as cv
import torch
import numpy as np
from .models.matching import Matching
from .models.utils import (frame2tensor)
torch.set_grad_enabled(False)


class SuperGlue:

    def __init__(self,  **kwargs):
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        print("run drive- " + str(self.device))
        self.keys = ['keypoints', 'scores', 'descriptors']
        config = {
            'superpoint': {
                'nms_radius': 5,
                'keypoint_threshold': 0.001,
                'max_keypoints': 4096
            },
            'superglue': {
                'weights': 'outdoor',
                'sinkhorn_iterations': 20,
                'match_threshold': 0.7,
            }
        }
        # Load SuperGlue model
        self.matcher = Matching(config).eval().to(self.device)

    def detect_features_matches(self, imgs, *args, **kwargs):
        features = []
        matche_infos = []
        for i in range(len(imgs)):

            img1 = cv.cvtColor(imgs[i], cv.COLOR_BGR2GRAY)
            frame_tensor_1 = frame2tensor(img1, self.device)
            data_1 = self.matcher.superpoint({'image': frame_tensor_1})
            data_1 = {k+'0': data_1[k] for k in self.keys}
            data_1['image0'] = frame_tensor_1
            size = (imgs[i].shape[1], imgs[i].shape[0])

            feature = self.convert_features(
                data_1['keypoints0'], data_1['scores0'], data_1['descriptors0'], size, i)
            features.append(feature)

            for j in range(len(imgs)):
                print(str(i)+"-"+str(j))
                matches_info = cv.detail_MatchesInfo()
                if i == j:
                    matches_info.src_img_idx = -1
                    matches_info.dst_img_idx = -1
                    matches_info.confidence = 0
                    matches_info.matches = ()
                    matches_info.inliers_mask = ()
                    matches_info.num_inliers = 0
                    matches_info.H = None
                else:
                    img2 = cv.cvtColor(imgs[j], cv.COLOR_BGR2GRAY)
                    frame_tensor_2 = frame2tensor(img2, self.device)
                    pred = self.matcher({**data_1, 'image1': frame_tensor_2})
                    kpts0 = data_1['keypoints0'][0].cpu().numpy()
                    kpts1 = pred['keypoints1'][0].cpu().numpy()
                    matches = pred['matches0'][0].cpu().numpy()
                    confidence = pred['matching_scores0'][0].detach(
                    ).cpu().numpy()

                    valid = matches > -1
                    mkpts0 = kpts0[valid]
                    mkpts1 = kpts1[matches[valid]]
                    
                    if len(mkpts0) > 10:
                        matches_info.src_img_idx = i
                        matches_info.dst_img_idx = j
                        if len(confidence[valid]) > 0:
                            matches_info.confidence = confidence[valid].mean()
                        else:
                            matches_info.confidence = 0
                        matches_info.matches = [cv.DMatch(
                            h, matches[h], confidence[h]) for h in range(len(matches)) if matches[h] > -1]
                        matches_info.inliers_mask = valid[valid]
                        matches_info.num_inliers = np.sum(valid)
                        matches_info.H, status = cv.findHomography(
                            mkpts0, mkpts1, cv.RANSAC, 5.0)
                    else:
                        matches_info.src_img_idx = i
                        matches_info.dst_img_idx = j
                        matches_info.confidence = 0
                        matches_info.matches = ()
                        matches_info.inliers_mask = ()
                        matches_info.num_inliers = 0
                        matches_info.H = None

                matche_infos.append(matches_info)

        return features, matche_infos

    def convert_features(self, keypoints, scores, desc, size, idx):
        img_features = cv.detail.ImageFeatures()
        img_features.img_idx = idx
        img_features.img_size = size
        img_features.descriptors = self.convert_descriptors(desc)
        img_features.keypoints = self.convert_keypoints(keypoints, scores)
        print('kp len')
        print(len(img_features.keypoints))

        return img_features

    def convert_keypoints(self, keypoints, scores):
        keypoints = keypoints[0].cpu().numpy()
        scores = scores[0].cpu().numpy()
        keypoints_cv = []

        for i in range(len(keypoints)):
            x = keypoints[i][0].item()
            y = keypoints[i][1].item()
            size = 10
            response = scores[i].item()
            octave = 0
            class_id = -1
            kp = cv.KeyPoint(x, y, size, response, octave, class_id)
            keypoints_cv.append(kp)

        return keypoints_cv

    def convert_descriptors(self, desc):
        descriptors = np.array([d.cpu().numpy() for d in desc])
        return cv.Mat(descriptors)