# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/12_consine.ipynb (unless otherwise specified).

__all__ = ['CosineSearch']

# Cell
import numpy as np
import pandas as pd

# Cell
class CosineSearch:
    """
    Build a index search on cosine distance
    cos = CosineSearch(base_array)
    idx_order = cos(vec)
    """

    def __init__(self, base: np.ndarray):
        """
        base: np.ndarray, embedding matrix of shape:
            (num_items, hidden_size)
        """
        assert len(base.shape) == 2,\
            f"Base array has to be 2 dimentional, input is {len(base.shape)}"
        self.base = base
        self.base_norm = self.calc_base_norm(self.base)
        self.normed_base = self.base/self.base_norm[:, None]
        self.dim = self.base.shape[1]

    def __repr__(self):
        return f"[Consine Similarity Search] ({len(self)} items)"

    def __len__(self): return self.base.shape[0]

    @staticmethod
    def calc_base_norm(base: np.ndarray) -> np.ndarray:
        return np.sqrt(np.power(base, 2).sum(1))

    def search(self, vec: np.ndarray, return_similarity: bool = False):
        if return_similarity:
            similarity = (vec * self.normed_base /
                          (np.power(vec, 2).sum())).sum(1)
            order = similarity.argsort()[::-1]
            return order, similarity[order]
        return self(vec)

    def __call__(self, vec: np.ndarray) -> np.ndarray:
        """
        Return the order index of the closest vector to the furthest
        vec: an 1 dimentional vector, marks the closest index
            to the further ones
        """
        return (vec * self.normed_base).sum(1).argsort()[::-1]