from jina import Executor, Document, DocumentArray, requests
import os
from utils.helper import clean_string
# import numpy as np
# from typing import Tuple

top_k = 1


class DiskIndexer(Executor):
    """Simple indexer class"""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._docs = DocumentArray()
        self.top_k = top_k
        if os.path.exists(self.save_path):
            self._docs = DocumentArray.load(self.save_path)
        else:
            self._docs = DocumentArray()

    @property
    def save_path(self):
        if not os.path.exists(self.workspace):
            os.makedirs(self.workspace)
        return os.path.join(self.workspace, "apps.json")

    def close(self):
        self._docs.save(self.save_path)

    @requests(on="/index")
    def index(self, docs: "DocumentArray", **kwargs):
        self._docs.extend(docs)
        return docs

    @requests(on="/search")
    def search(self, docs, **kwargs):
        print(
            f"\tSearching index of {len(self._docs)} Documents for \"{docs[0].text}\"")
        darr = self._docs
        apis = darr.get_attributes('text')
        scores = {}
        for api in apis:
            api = clean_string(api)
            scores[api] = 0
        search_term = docs[0].text

        s_list = clean_string(search_term).split()

        for x in range(len(scores)):
            api_desc = list(scores)[x]
            for word in s_list:
                if word in api_desc:
                    if len(word) > 2 and word not in "and the which was were whose has have been":
                        scores[api_desc] += 1

        a = sorted(scores.items(), key=lambda x: x[1])
        result = a[-1][0]
        required_doc = Document()

        for x in range(len(apis)):
            api = clean_string(apis[x])
            if api == result:
                required_doc = darr[x]

        return DocumentArray([required_doc])


#     @requests(on="/search")
#     def search(self, docs: "DocumentArray", **kwargs):
#         print(
#             f"\tSearching index of {len(self._docs)} Documents for \"{docs[0].text}\"")
#         a = np.stack(docs.get_attributes("embedding"))
#         b = np.stack(self._docs.get_attributes("embedding"))
#         q_emb = _ext_A(_norm(a))
#         d_emb = _ext_B(_norm(b))
#         dists = _cosine(q_emb, d_emb)
#         idx, dist = self._get_sorted_top_k(dists, self.top_k)
#         for _q, _ids, _dists in zip(docs, idx, dist):
#             for _id, _dist in zip(_ids, _dists):
#                 d = Document(self._docs[int(_id)], copy=True)
#                 d.score.value = 1 - _dist
#                 _q.matches.append(d)
#         return docs


#     @staticmethod
#     def _get_sorted_top_k(
#         dist: "np.array", top_k: int
#     ) -> Tuple["np.ndarray", "np.ndarray"]:
#         if top_k >= dist.shape[1]:
#             idx = dist.argsort(axis=1)[:, :top_k]
#             dist = np.take_along_axis(dist, idx, axis=1)
#         else:
#             idx_ps = dist.argpartition(kth=top_k, axis=1)[:, :top_k]
#             dist = np.take_along_axis(dist, idx_ps, axis=1)
#             idx_fs = dist.argsort(axis=1)
#             idx = np.take_along_axis(idx_ps, idx_fs, axis=1)
#             dist = np.take_along_axis(dist, idx_fs, axis=1)

#         return idx, dist


# def _get_ones(x, y):
#     return np.ones((x, y))


# def _ext_A(A):
#     nA, dim = A.shape
#     A_ext = _get_ones(nA, dim * 3)
#     A_ext[:, dim: 2 * dim] = A
#     A_ext[:, 2 * dim:] = A ** 2
#     return A_ext


# def _ext_B(B):
#     nB, dim = B.shape
#     B_ext = _get_ones(dim * 3, nB)
#     B_ext[:dim] = (B ** 2).T
#     B_ext[dim: 2 * dim] = -2.0 * B.T
#     del B
#     return B_ext


# def _euclidean(A_ext, B_ext):
#     sqdist = A_ext.dot(B_ext).clip(min=0)
#     return np.sqrt(sqdist)


# def _norm(A):
#     return A / np.linalg.norm(A, ord=2, axis=1, keepdims=True)


# def _cosine(A_norm_ext, B_norm_ext):
#     return A_norm_ext.dot(B_norm_ext).clip(min=0) / 2
