# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/06_indexing.ipynb.

# %% auto 0
__all__ = ['context', 'store', 'VectorNodesIndex']

# %% ../nbs/06_indexing.ipynb 2
from .base import *
from .store import *
from .context import *
from .llm import *
from .loaders import *
from typing import Union, List, Dict, Tuple, Optional, Any
import numpy as np
from abc import ABC, abstractmethod

# %% ../nbs/06_indexing.ipynb 3
#| eval: false
context = ModelContext()
context.set_default()
store = DocumentStore()

# %% ../nbs/06_indexing.ipynb 7
class VectorNodesIndex: #Compatible with TextNode right now. Storage of reference of certain nodes.
    #Try out with SVM retrieval strategy and other ones. 
    #Question. Treat docstore the
    #TODO: Proper context validation
    """Inside here the embeddings stored are normalized. So when doing operations with vectors has to be kept into account. """
    def __init__(self, context): #May not be needed in postgres. 
        self.idx_to_node = {}
        self.idx = np.array([], dtype=np.int64)
        self.context = context
        #This line below accepts huggingface embeddings format. 
        self.embedding_dim = self.context.embedding[1].word_embedding_dimension
        if self.embedding_dim is not None:
            self.embeddings = np.empty((0, self.embedding_dim))
        else:
            self.embeddings = np.array([])
        self.retrieval_strategies = {
            "dot_product": self.retrieve_dot_product,
        }

    def add(self, nodes: Union[TextNode, List[TextNode]]): #Embed with non excluded content. 
        ##TODO: Add nodes to docstore that is referenced there by default. 
        # And retrive them directly
        if isinstance(nodes, TextNode):
            nodes = [nodes]
        elif isinstance(nodes, list):
            new_embeddings = np.vstack([node.embedding for node in nodes])
            if self.embeddings.size == 0:
                self.embeddings = new_embeddings
            else:
                self.embeddings = np.append(self.embeddings, new_embeddings, axis=0)
            node_idx = np.arange(len(self.idx), self.embeddings.shape[0])
            for node, idx in zip(nodes, node_idx):
                self.idx_to_node[idx] = node.id
                node.idx_ref = idx
            self.idx = np.concatenate((self.idx, node_idx))

    def get_node_ids(self, idx_refs: List[int]):
        """Providing a list of idx_refs of embeddings of the nodes, get the corresponding node ids"""
        return [self.idx_to_node[idx_ref] for idx_ref in idx_refs]

    def get_embedding(self, idx_ref: Union[List[int], int]):
        """Providing the idx_ref of the node, or nodes get the embedding"""
        if isinstance(idx_ref, np.int64) or isinstance(idx_ref, int):
            idx_ref = np.array([idx_ref], dtype = np.issubdtype(type(idx_ref), np.integer))
        if isinstance(idx_ref, list):
            idx_ref = np.array(idx_ref, dtype= np.int64)
        if isinstance(idx_ref, np.ndarray):
            return self.embeddings[idx_ref]

    def retrieve(self, query_str: List[str], top_k: int = 10, strategy: str = 'dot_product'):
        query_embedding = self.context.embedding.encode(query_str)
        query_embedding = query_embedding / np.linalg.norm(query_embedding)  # Ensure normalization
        retrieve_method = self.retrieval_strategies.get(strategy, self.retrieve_dot_product)
        top_k_idx_matches, top_k_scores = retrieve_method(query_embedding, top_k)
        nodes = self.get_node_ids(top_k_idx_matches)
        return nodes, top_k_scores
    
    def retrieve_dot_product(self, query_embedding, top_k):
        similarity = self.cosine_similarity(query_embedding, self.embeddings.T)[0]
        top_matches_similarity = np.argsort(similarity)[::-1]
        top_k_idxs = top_matches_similarity[:top_k]
        top_k_scores = similarity[top_k_idxs]
        return top_k_idxs, top_k_scores
    
    def cosine_similarity(self,vecA, vecB):
        return np.dot(vecA, vecB)
