"""Base VectorStore class"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/01_ingest.stores.base.ipynb.

# %% auto 0
__all__ = ['VectorStore']

# %% ../../../nbs/01_ingest.stores.base.ipynb 3
from abc import ABC, abstractmethod
import os
from typing import Optional, Callable
from langchain_huggingface import HuggingFaceEmbeddings
from ..base import DEFAULT_CHUNK_SIZE, DEFAULT_CHUNK_OVERLAP, TABLE_CHUNK_SIZE, CHROMA_MAX, process_folder


class VectorStore(ABC):

    def ingest(
        self,
        source_directory: str, # path to folder containing document store
        chunk_size: int = DEFAULT_CHUNK_SIZE, # text is split to this many characters by [langchain.text_splitter.RecursiveCharacterTextSplitter](https://api.python.langchain.com/en/latest/character/langchain_text_splitters.character.RecursiveCharacterTextSplitter.html)
        chunk_overlap: int = DEFAULT_CHUNK_OVERLAP, # character overlap between chunks in `langchain.text_splitter.RecursiveCharacterTextSplitter`
        ignore_fn:Optional[Callable] = None, # Optional function that accepts the file path (including file name) as input and returns `True` if file path should not be ingested.
        batch_size:int=CHROMA_MAX, # batch size used when processing documents
        **kwargs
    ) -> None:
        """
        Ingests all documents in `source_directory` (previously-ingested documents are
        ignored). When retrieved, the
        [Document](https://api.python.langchain.com/en/latest/documents/langchain_core.documents.base.Document.html)
        objects will each have a `metadata` dict with the absolute path to the file
        in `metadata["source"]`.
        Extra kwargs fed to `ingest.load_single_document`.
        """

        if not os.path.exists(source_directory):
            raise ValueError("The source_directory does not exist.")
        elif os.path.isfile(source_directory):
            raise ValueError(
                "The source_directory argument must be a folder, not a file."
            )
        texts = None
        if self.exists():
            # Update and store locally vectorstore
            print(f"Appending to existing vectorstore at {self.persist_location}")
            # Get existing sources and normalize them for consistent comparison
            ignored_files = set([os.path.normpath(d['source']).replace('\\', '/') 
                                for d in self.get_all_docs()])
        else:
            print(f"Creating new vectorstore at {self.persist_location}")
            ignored_files = []

        texts = process_folder(
            source_directory,
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            ignored_files=ignored_files,
            ignore_fn=ignore_fn,
            batch_size=batch_size,
            **kwargs

        )

        texts = list(texts)
        print(f"Split into {len(texts)} chunks of text (max. {chunk_size} chars each for text; max. {TABLE_CHUNK_SIZE} chars for tables)")

        self.add_documents(texts, batch_size=batch_size)

        if texts:
            print(
                "Ingestion complete! You can now query your documents using the LLM.ask or LLM.chat methods"
            )
        db = None
        return
    
    def init_embedding_model(self, 
                             embedding_model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
                             embedding_model_kwargs: Optional[dict] = None,
                             embedding_encode_kwargs: dict = {"normalize_embeddings": False},
                             **kwargs
                             ):
        """
        Instantiate embedding model
        """
        if not embedding_model_kwargs:
            import torch
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            embedding_model_kwargs = {"device": device}          
        self.embeddings =  HuggingFaceEmbeddings(model_name=embedding_model_name, 
                                     model_kwargs=embedding_model_kwargs,
                                     encode_kwargs=embedding_encode_kwargs)

    def get_embedding_model(self):
        """
        Returns an instance to the `langchain_huggingface.HuggingFaceEmbeddings` instance
        """
        return self.embeddings


    def compute_similarity(self, query:str, texts:list):
        """
        Computes semantic similarity between a query and a list of texts
        """
        from sentence_transformers import util
        import torch

        embeddings = self.get_embedding_model()

        # Compute embeddings
        query_emb = torch.tensor(embeddings.embed_query(query)).unsqueeze(0)  # Shape (1, embedding_dim)
        text_embs = torch.tensor(embeddings.embed_documents(texts))  # Shape (len(texts), embedding_dim)

        # Compute cosine similarity
        cos_scores = util.pytorch_cos_sim(query_emb, text_embs).squeeze(0).tolist()  # Shape (len(texts),)

        return cos_scores


    def check(self):
        """
        Raise exception if `VectorStore.exists()` returns False
        """
        if not self.exists():
            raise Exception('The vector store is either empty or does not yet exist. '+\
                            'Please invoke the `ingest` method or `add_document` method.')


    def normalize_text(self, text):
        """
        normalize text (e.g., from "classiﬁcation" to "classification")
        """
        import unicodedata
        try:
            return unicodedata.normalize('NFKC', text)
        except:
            return text
 

    @abstractmethod
    def exists(self):
        """
        Returns True if vector store has been initialized and contains documents.
        """
        pass


    @abstractmethod
    def add_documents(self, documents):
        """
        Stores instances of `langchain_core.documents.base.Document` in vector store
        """
        pass
        
    
    @abstractmethod
    def remove_document(self, id_to_delete):
        """
        Remove a single document.
        """
        pass

    @abstractmethod
    def remove_source(self, source):
        """
        Remove documents by source
        """
        pass
    
    @abstractmethod
    def update_documents(self, *args, **kwargs):
        """
        Update a set of documents.
        """
        pass

    @abstractmethod
    def get_all_docs(self):
        """
        Returns a list of files previously added to vector store.
        """
        pass


    @abstractmethod
    def get_size(self):
        """
        Get total number of records added to vector store
        """
        pass

    @abstractmethod
    def erase(self):
        """
        Removes all documents in vector store
        """
        pass


    @abstractmethod
    def get_doc(self, id):
        """
        Retrieve a document by ID
        """
        pass


    @abstractmethod
    def search(self, query, limit=4):
        """
        Searches the vector store.
        For sparse stores, this is simply a keyword-search.
        For dense stores, this is equivalent to semantic_search except results
        are in the form of dictionary with keys 'hits' and 'total_hits'.
        Method must include 'query' as first positional argument and "limit" as keyword argument.

        """
        pass

    @abstractmethod
    def semantic_search(self):
        """
        Semantic search of vector store.
        Expected to return a list of Langchain Document objects.
        """
        pass

    def query(self, query: str, **kwargs):
        """
        Generic query method that invokes the store's search method.
        This provides a consistent interface across all store types.
        """
        return self.search(query, **kwargs)
