import bm25s
from collections import defaultdict
import pandas as pd
from tangbao.apollo import Apollo
from tangbao.parse_docs import separate_for_BM25
from tangbao.parse_docs import separate_for_indexing
from langchain_core.vectorstores import InMemoryVectorStore

# Define the BM25Retriever class, a wrapper around the bm25s library for BM25 (keyword-based) retrieval.
class BM25Retriever:
    def __init__(self, corpus: list):
        """
        Initialize the BM25Retriever with a given corpus.

        Parameters:
        corpus (list): A list of dictionaries, containing document texts (strings) and metadatas (dicts) to be indexed.
        """
        self.corpus = corpus
        self.retriever = bm25s.BM25(corpus=self.corpus)
        self.corpus_tokens = None

    def index_corpus(self) -> None:
        """
        Tokenize and index the corpus.
        """
        # Extract the text from the documents
        corpus_text = [doc["text"] for doc in self.corpus]
        # Tokenize the corpus and only keep the ids (faster and saves memory)
        self.corpus_tokens = bm25s.tokenize(corpus_text, stopwords="en")
        # Index the corpus
        self.retriever.index(self.corpus_tokens)

    def query_corpus(self, query: str, k: int = 2) -> tuple:
        """
        Query the indexed corpus and retrieve the top-k results.

        Parameters:
        query (str): The query string.
        k (int): The number of top results to retrieve. Defaults to 2 because its usually used in tandem with other retrievers.

        Returns:
        output: A list of dictionaries containing the metadata, text, and score of the top-k results.
        """
        # Tokenize the query
        query_tokens = bm25s.tokenize(query)
        # Get top-k results as a tuple of (doc ids, scores). Both are arrays of shape (n_queries, k)
        results, scores = self.retriever.retrieve(query_tokens, corpus=self.corpus, k=k)
        # Store results in a list
        output = []
        for idx in range(results.shape[1]):
            if scores[0, idx] != 0: # do not include documents with score 0 (no match)
                output.append(
                    {
                        "metadata": results[0, idx]["metadata"],
                        "text": results[0, idx]["text"],
                        "score": scores[0, idx],
                    }
                )

        return output
    

def extract_docs_from_apollo_vs(response: dict) -> list:
    """
    Extracts text, metadata and score from a response generated by apollo.query_index.

    Args:
        response (dict): A dictionary containing a list of documents retrieved from apollo.query_index.
                         Each document is expected to have 'text', 'metadata' and 'scores' keys.

    Returns:
        list: A list of dictionaries, each containing 'text', 'metadata' and 'score' from the documents in the response.
    """
    extraction = [{"text": doc['text'], "metadata": doc['metadata'], "score": doc['score']} for doc in response['docs']]
    return extraction

def extract_docs_from_langchain_re(response: dict) -> list:
    """
    Extracts text, metadata and score from a response generated by a langchain retriever.

    Args:
        response (dict): A dictionary containing a list of langchain Documents retrieved from a langchain retriever.
                         Each Document is expected to have 'page_content' and 'metadata' keys.

    Returns:
        list: A list of dictionaries, each containing 'text' and 'metadata' from the documents in the response.
    """
    extraction = [{"text": doc.page_content, "metadata": doc.metadata} for doc in response]
    return extraction

def rrf(rank_lists, k=60, weights=None):
    """
    Perform Reciprocal Rank Fusion on multiple ranked lists.

    :param rank_lists: List of ranked lists. Each ranked list is a list of dicts with "text" and "metadata" keys.
    :param k: Constant to control the smoothness of the scores. Defaults to 60 based on the original RRF paper by Cormack et al. (2009).
    :param weights: Weights for the scores of the ranked lists.
    :return: A list of dicts with keys "rank", "text" and "metadata".
    """

    # Initialize a dictionary to store the combined scores
    combined_scores = defaultdict(float)
    metadata_dict = {}

    # If weights are not provided, set them to 1 (standard RRF)
    if weights is None:
        weights = [1] * len(rank_lists)

    # Iterate over each ranked list
    for id_list, rank_list in enumerate(rank_lists):
        # Iterate over each document in the ranked list
        for rank, doc in enumerate(rank_list):
            doc_id = doc['text']
            metadata_dict[doc_id] = doc['metadata']
            # Calculate the RRF score for the document
            score = weights[id_list] / (k + rank + 1)
            # Add the score to the combined score for the document
            combined_scores[doc_id] += score

    # Create a list of dicts with "rank", "text" and "metadata"
    sorted_docs = sorted(combined_scores.items(), key=lambda item: item[1], reverse=True)
    result = [{'rank': score, 'text': doc_id, 'metadata': metadata_dict[doc_id]} for doc_id, score in sorted_docs if score > 0]

    return result


class HybridRetrieverApollo:
    """
    A hybrid retriever that combines BM25 and dense retrieval methods using Reciprocal Rank Fusion (RRF).
    The dense retrieval is performed using the apollo vector search engine.
    NOTE: this class requires the apollo vector search engine to be previously populated with documents.

    Attributes:
        df (pd.DataFrame): The dataframe containing the documents.
        index_name (str): The name of the index to be used by the apollo retriever.
        embedding_model (str): The name of the embedding model to be used by the apollo retriever.
        documents (list): A list of documents separated for BM25 retrieval.
        bm25_retriever (BM25Retriever): An instance of the BM25Retriever class.
        apollo (Apollo): An instance of the apollo class for dense retrieval.

    Methods:
        query_indices(user_query: str, num_chunks: int, keyword_weight: float=None) -> list:
            Queries both BM25 and dense retrievers and combines the results using RRF.
            Args:
                user_query (str): The query string provided by the user.
                num_chunks (int): The number of chunks to retrieve.
                keyword_weight (float, optional): The weight to assign to the BM25 results in the RRF. Defaults to None. Can range between 0 and 1.
            Returns:
                list: A list of combined scores from both retrievers.
    """
    def __init__(self, df: pd.DataFrame, index_name: str, embedding_model: str):
        self.df = df
        self.index_name = index_name
        self.embedding_model = embedding_model
        self.documents = separate_for_BM25(df)
        self.bm25_retriever = BM25Retriever(self.documents)
        self.bm25_retriever.index_corpus()
        self.apollo = Apollo()

    def query_indices(self, user_query: str, num_chunks: int, keyword_weight: float=None):
        # fetch results from the BM25 retriever
        bm25_docs = self.bm25_retriever.query_corpus(query=user_query, k=num_chunks)
        
        # fetch results from the apollo retriever
        response_apollo_query_index = self.apollo.query_index(user_query=user_query, num_chunks=num_chunks, index_name=self.index_name, embedding_model=self.embedding_model)
        
        # format the apollo results in the same way as the BM25 results
        dense_docs = extract_docs_from_apollo_vs(response_apollo_query_index)
        
        # Modify the input to be a list of dicts
        rank_lists = [bm25_docs, dense_docs]
        
        # Perform RRF
        if keyword_weight:
            combined_scores = rrf(rank_lists, weights=[keyword_weight, 1-keyword_weight])
        else: 
            combined_scores = rrf(rank_lists)
        
        return combined_scores

class ApolloEmbeddings:
    def __init__(self, model: str = "openai-text-embedding-3-small", dimensions: int = 512):
        """
        Initialize the ApolloEmbeddings class. This class is used to generate embeddings using the Apollo API and it is introduced to facilitate the integration with LangChain retrievers.

        Args:
            model (str): The model to use for generating embeddings. Default is "openai-text-embedding-3-small".
            dimensions (int): The number of dimensions for the embeddings. Default is 512. NOTE: only openai models support this parameter. The other models will just ignore it.
        """
        self.apollo: Apollo = Apollo()
        self.model: str = model
        self.dimensions: int = dimensions

    def embed_query(self, text: str) -> list:
        """
        Embed a single query.

        Args:
            text (str): A string to be embedded.

        Returns:
            list: An array of embeddings.
        """
        response_full = self.apollo.get_embeddings(model=self.model, texts=[text], dimensions=self.dimensions)
        response = [item.embedding for item in response_full.data]
        return response[0]

    def embed_documents(self, texts: list) -> list:
        """
        Embed multiple documents at once.

        Args:
            texts (list): A list of strings to be embedded.

        Returns:
            list: A list of embeddings.
        """
        response_full = self.apollo.get_embeddings(model=self.model, texts=texts, dimensions=self.dimensions)
        response = [item.embedding for item in response_full.data]
        return response

class HybridRetrieverInMemory:
    """
    A hybrid retriever that combines BM25 and dense retrieval methods using Reciprocal Rank Fusion (RRF).
    The dense retrieval is performed using a LangChain retriever.
    
    Attributes:
        df (pd.DataFrame): The dataframe containing the documents to be indexed.
        documents (list): A list of documents separated for BM25 retrieval.
        texts (list): A list of text chunks to be indexed.
        ids (list): A list of ids to be indexed.
        metadatas (list): A list of metadatas to be indexed.
        bm25_retriever (BM25Retriever): An instance of the BM25Retriever class.
        embeddings (ApolloEmbeddings): An instance of the ApolloEmbeddings class.
        langchain_vectorstore (VectorStore): An instance of a LangChain InMemoryVectorStore object, populated with the documents in the dataframe.

    Methods:
        query_indices(user_query: str, num_chunks: int, keyword_weight: float=None) -> list:
            Queries both BM25 and dense retrievers and combines the results using RRF.
            Args:
                user_query (str): The query string provided by the user.
                num_chunks (int): The number of chunks to retrieve.
                keyword_weight (float, optional): The weight to assign to the BM25 results in the RRF. Defaults to None. Can range between 0 and 1.
            Returns:
                list: A list of combined scores from both retrievers.
    """
    def __init__(self, df: pd.DataFrame, emb_model: str = "openai-text-embedding-3-small", emb_dimensions: int = 512):
        self.df = df
        self.documents = separate_for_BM25(df)
        self.texts, self.ids, self.metadatas = separate_for_indexing(df)
        self.bm25_retriever = BM25Retriever(self.documents)
        self.bm25_retriever.index_corpus()
        self.embeddings = ApolloEmbeddings(model=emb_model, dimensions=emb_dimensions)
        self.langchain_vectorstore = InMemoryVectorStore.from_texts(
                                        texts=self.texts,
                                        embedding=self.embeddings,
                                        ids=self.ids,
                                        metadatas=self.metadatas,
                                    )

    def query_indices(self, user_query: str, num_chunks: int, keyword_weight: float=None):
        # fetch results from the BM25 retriever
        bm25_docs = self.bm25_retriever.query_corpus(query=user_query, k=num_chunks)
        
        # fetch results from the langchain in memory retriever
        self.retriever = self.langchain_vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": num_chunks})
        response_inmemory_retriever = self.retriever.invoke(user_query)

        # format the apollo results in the same way as the BM25 results
        dense_docs = extract_docs_from_langchain_re(response_inmemory_retriever)
        
        # Modify the input to be a list of dicts
        rank_lists = [bm25_docs, dense_docs]
        
        # Perform RRF
        if keyword_weight:
            combined_scores = rrf(rank_lists, weights=[keyword_weight, 1-keyword_weight])
        else: 
            combined_scores = rrf(rank_lists)
        
        return combined_scores