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

__all__ = ['CosineSearch', 'InterpEmbeddings']

# Cell
import numpy as np
import pandas as pd
from typing import Dict, List, Any
from forgebox.html import DOM

# 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]

# Cell
class InterpEmbeddings:
    """
    interp = InterpEmbeddings(embedding_matrix, vocab_dict)

    interp.search("computer")

    # visualize the embedding with tensorboard
    interp.visualize_in_tb()
    """

    def __init__(
        self,
        embedding_matrix: np.ndarray,
        vocab: Dict[int, str]
    ):
        """
        embedding_matrix: np.ndarray, embedding matrix of shape:
            (num_items, hidden_size)
        """
        self.base = embedding_matrix
        self.cosine = CosineSearch(embedding_matrix)
        self.vocab = vocab
        self.c2i = dict((v, k) for k, v in vocab.items())

    def __repr__(self) -> str:
        cls = self.__class__.__name__
        return f"{cls} with\n\t{self.cosine}"

    def search(
        self,
        category: str,
        top_k: int = 20,
    ) -> pd.DataFrame:
        """
        search for similar words with embedding and vocabulary dictionary
        """
        token_id = self.c2i.get(category)
        if token_id is None:
            match_list = []
            for token_id, token in self.vocab.items():
                if category.lower() in str(token).lower():
                    match_list.append({"token": token, "token_id": token_id})
            if len(match_list)==0:
                raise KeyError(
                    f"[UnpackAI] category: {category} not in vocabulary")
            else:
                match_df = pd.DataFrame(match_list)
                DOM("Search with the following categories","h3")()
                display(match_df)
                token_ids = list(match_df.token_id)
        else:
            DOM(f"Search with token id {token_id}","h3")()
            token_ids = [token_id,]

        # combine multiple tokens into 1
        vec = self.base[token_ids].mean(0)

        # distance search
        closest, similarity = self.cosine.search(vec, return_similarity=True)

        closest = closest[:top_k]
        similarity = similarity[:top_k]
        tokens = list(self.vocab.get(idx) for idx in closest)
        return pd.DataFrame({
            "tokens": tokens,
            "idx": closest,
            "similarity": similarity})

    def visualize_in_tb(
        self,
        log_dir:str="./logs",
        selection: np.ndarray=None,
        first_k:int=500,
    ) -> None:
        """
        Visualize the embedding in tensorboard
        For now this function is only supported on colab
        """
        # since this won't be excute too many times within a notebook
        # in large chances... so to avoid missing library when import
        # other function under this module: we import related stuff here
        from torch.utils.tensorboard import SummaryWriter
        # this version's pd has vc for quick value counts
        from forgebox.imports import pd
        import tensorflow as tf
        import tensorboard as tb
        import os

        # possible tensorflow version error
        tf.io.gfile = tb.compat.tensorflow_stub.io.gfile
        os.system(f"rm -rf {log_dir}")
        writer = SummaryWriter(log_dir=log_dir,)
        self.i2c = dict((v,k) for k,v in self.c2i.items())
        tokens = list(self.i2c.get(i) for i in range(len(self.i2c)))

        if selection is None:
            vecs = self.base[:first_k]
            tokens = tokens[:first_k]
        else:
            selection = np.array(selection).astype(dtype=np.int64)
            # select a pool of tokens for visualizaiton
            tokens = list(np.array(tokens)[selection][:first_k])
            vecs = self.base[selection][:first_k]
        writer.add_embedding(vecs, metadata=tokens,)
        # prompts for next step
        print(f"Please run the the following command in a cell")
        print("%load_ext tensorboard")
        print(f"%tensorboard  --logdir {log_dir}")