from typing import List
import math

from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever

class BM25Retriever(BaseRetriever):
    documents: List[Document]
    k: int
    document_entity_index: dict
    idf_scores: dict
    average_doc_length: float

    def __init__(self, documents: List[Document], k: int, entity_extractor):
        super().__init__()
        self.documents = documents
        self.k = k
        self.entity_extractor = entity_extractor
        self._prepare_documents()

    def _prepare_documents(self):
        """Prepare documents by extracting entities and calculating necessary BM25 metrics."""
        num_docs = len(self.documents)
        doc_lengths = []
        df = {}
        
        # Extract entities and calculate document frequency (DF)
        for doc in self.documents:
            entities = self.entity_extractor.extract(doc.page_content)
            doc.entity_bag = entities
            doc_lengths.append(len(entities))
            
            unique_entities = set(entities)
            for entity in unique_entities:
                if entity in df:
                    df[entity] += 1
                else:
                    df[entity] = 1
        
        self.average_doc_length = sum(doc_lengths) / num_docs
        self.idf_scores = {term: math.log((num_docs - df[term] + 0.5) / (df[term] + 0.5)) for term in df}

    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        """Use BM25 to rank documents based on entities extracted from the query."""
        query_entities = self.entity_extractor.extract(query)
        scores = []

        for doc in self.documents:
            score = 0
            for entity in query_entities:
                if entity in doc.entity_bag:
                    term_frequency = doc.entity_bag.count(entity)
                    idf = self.idf_scores.get(entity, 0)
                    doc_length = len(doc.entity_bag)
                    score += idf * (term_frequency * (1.2 + 1) / (term_frequency + 1.2 * (1 - 0.75 + 0.75 * (doc_length / self.average_doc_length))))
            if score > 0:
                scores.append((score, doc))

        # Sort documents by their score and return top k
        sorted_docs = sorted(scores, key=lambda x: x[0], reverse=True)
        return [doc for _, doc in sorted_docs[:self.k]]

# Note: This implementation assumes the presence of an entity_extractor with an extract method.