import cv2
import numpy as np

import LFUtilities
import BEBLIDParameters
import WebAppSettings as settings


class BEBLIDSearcher:

    def __init__(self):
        self.lf = LFUtilities.load(settings.DATASET_BEBLID)
        self.ids = np.loadtxt(settings.DATASET_IDS, dtype=str).tolist()
        #self.bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
        self.bf = cv2.DescriptorMatcher_create(cv2.DescriptorMatcher_BRUTEFORCE_HAMMING)

    def get_id(self, idx):
        return self.ids[idx]

    def search_by_id(self, query_id):
        query_idx = self.ids.index(query_id)
        return self.search_by_img(self.lf[query_idx])

    def search_by_img(self, query):
        max_inliers = -1
        res = []

        for img_id, features in zip(self.ids, self.lf):
            try:
                nn_matches = self.bf.knnMatch(query[1], features[1], 2)
                good = [m for m, n in nn_matches if m.distance < BEBLIDParameters.NN_MATCH_RATIO * n.distance]

                if len(good) > BEBLIDParameters.MIN_GOOD_MATCHES:
                    src_pts = np.float32([query[0][m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
                    dst_pts = np.float32([features[0][m.trainIdx].pt for m in good]).reshape(-1, 1, 2)

                    M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 1.0)
                    matches_mask = mask.ravel().tolist()
                    # print(len(good))
                    inliers = np.count_nonzero(matches_mask)
                    # print(inliers)
                    if (inliers >= BEBLIDParameters.MIN_INLIERS and inliers > max_inliers):
                        max_inliers = inliers
                        class_name = img_id.split('/')[0]
                        res.append((class_name, round(inliers/len(good), 3)))
            except:
                print('search error evaluating ' + img_id)
                pass

        if res:
            res.sort(key=lambda result: result[1], reverse=True)
            return res

    def add(self, lf):
        self.lf.append(lf)

    def remove(self, idx):
        self.descs = np.delete(self.descs, idx, axis=0)

    def save(self, is_backup=False):
        lf_save_file = settings.DATASET_LF
        ids_file = settings.DATASET_IDS_LF
        if lf_save_file != "None":
            if is_backup:
                lf_save_file += '.bak'
                ids_file += '.bak'

            LFUtilities.save(lf_save_file, self.lf)
            np.savetxt(ids_file, self.ids, fmt='%s')