import pickle
import faiss
import numpy as np
# from grammar import remove_verbs, clean_text
from utils import *
from sentence_transformers import SentenceTransformer


class FAISS:
    def __init__(self, dimensions: int):
        self.dimensions = dimensions
        self.index = faiss.IndexFlatL2(dimensions)
        self.vectors = {}
        self.counter = 0
        self.model_name = 'paraphrase-multilingual-MiniLM-L12-v2'
        self.sentence_encoder = SentenceTransformer(self.model_name)

    def init_vectors(self, path):
        with open(path, 'rb') as pkl_file:
            self.vectors = pickle.load(pkl_file)

    def init_index(self, path):
        self.index = faiss.read_index(path)

    def add(self, text, idx, pop, emb=None):
        if emb is None:
            text_vec = self.sentence_encoder.encode([text])
        else:
            text_vec = emb
        self.index.add(text_vec)
        self.vectors[self.counter] = (idx, text, pop, text_vec)
        self.counter += 1

    def search(self, v: list, k: int = 10):
        result = []
        distance, item_index = self.index.search(v, k)
        for dist, i in zip(distance[0], item_index[0]):
            if i == -1:
                break
            else:
                result.append((self.vectors[i][0], self.vectors[i][1], self.vectors[i][2], dist))

        return result

    def suggest_tags(self, query, top_n=10, k=30) -> list:

        emb = self.sentence_encoder.encode([query.lower()])
        r = self.search(emb, k)

        result = []
        for i in r:
            if check(query, i[1]):
                result.append(i)
        # надо добавить вес относительно длины
        result = sorted(result, key=lambda x: x[0] * 0.3 - x[-1], reverse=True)
        total_result = []
        for i in range(len(result)):
            flag = True
            for j in result[i + 1:]:
                flag &= sweet_check(result[i][1], j[1])
            if flag:
                total_result.append(result[i][1])

        return total_result[:top_n]