import torch
from transformers import pipeline
from typing import List
from rankify.models.base import BaseRanking
from rankify.dataset.dataset import Document
from tqdm import tqdm  # Import tqdm for progress tracking


class EchoRankReranker(BaseRanking):
    """
    A reranking model implementing the EchoRank method `[5]`_ for budget-constrained text reranking
    using large language models (LLMs). It employs a two-stage process:
    
    .. _[5]: https://arxiv.org/abs/2402.10866

    1. Binary Classification Stage: Determines if each passage is relevant to the query.
    2. Pairwise Ranking Stage: Ranks relevant passages via direct comparisons.

    EchoRank efficiently balances reranking performance and computational cost by limiting token usage.

    References
    ----------
    .. [5] Rashid, Muhammad Shihab, Jannat Ara Meem, Yue Dong, and Vagelis Hristidis. "EcoRank: Budget-Constrained Text Re-ranking Using Large Language Models." arXiv preprint arXiv:2402.10866 (2024).
    
    Attributes
    ----------
    method : str, optional
        The reranking method name.
    model_name : str
        The pretrained model name (default: `"google/flan-t5-large"`).
    type : str
        Type of budget constraint to apply (default: `"cheap"`).
    budget_tokens : int
        Total token budget for reranking (default: `4000`).
    budget_split_x : float
        Percentage of budget allocated to the binary classification stage (default: `0.5`).
    budget_split_y : float
        Percentage of budget allocated to the pairwise ranking stage (default: `0.5`).
    total_passages : int
        Maximum number of passages to process per query (default: `50`).
    device : str
        The device on which the model runs (`"cuda"` if available, otherwise `"cpu"`).
    model : transformers.Pipeline
        The Hugging Face pipeline for text-to-text generation used for reranking.

    See Also
    --------
    Reranking : Main interface for reranking models, including `EchoRankReranker`.

    Examples
    --------
    Basic usage with the `Reranking` interface:

    >>> from rankify.dataset.dataset import Document, Question, Answer, Context
    >>> from rankify.models.reranking import Reranking
    >>>
    >>> question = Question("What is climate change?")
    >>> answers = Answer(["Climate change refers to long-term shifts in temperatures and weather patterns."])
    >>> contexts = [
    >>>     Context(text="Climate change is mainly caused by human activities.", id=1),
    >>>     Context(text="Deforestation contributes to global warming.", id=2),
    >>>     Context(text="Polar bears are affected by melting ice caps.", id=3),
    >>>     Context(text="The stock market fluctuates daily.", id=4),
    >>> ]
    >>> document = Document(question=question, answers=answers, contexts=contexts)
    >>>
    >>> # Initialize Reranking with EchoRankReranker
    >>> model = Reranking(method='echorank', model_name='flan-t5-large')
    >>> model.rank([document])
    >>>
    >>> # Print reordered contexts
    >>> print("Reordered Contexts:")
    >>> for context in document.reorder_contexts:
    >>>     print(context.text)

    Notes
    -----
    - EchoRank is designed for budget-constrained ranking by limiting token usage.
    - Integrates into the `Reranking` class, so use `Reranking` instead of `EchoRankReranker` directly.
    - Uses two-stage ranking: binary filtering followed by pairwise comparisons.
    - This class implement from the following link https://github.com/shihabrashid-ucr/EcoRank/tree/main.
    """
    def __init__(self, method=None, model_name=None, **kwargs):
        """
        Initializes the EchoRank reranker.

        Parameters
        ----------
        method : str, optional
            The reranking method name.
        model_name : str, optional
            The Hugging Face model name to use for reranking (default: `"google/flan-t5-large"`).
        kwargs : dict
            Additional parameters for token budget and model configurations.
        """
        self.method = method
        self.model_name = model_name or kwargs.get("model_name", "google/flan-t5-large")
        self.type = kwargs.get("type", "cheap")
        self.budget_tokens = kwargs.get("budget_tokens", 4000)
        self.budget_split_x = kwargs.get("budget_split_x", 0.5)
        self.budget_split_y = kwargs.get("budget_split_y", 0.5)
        self.total_passages = kwargs.get("total_passages", 50)
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        # Load the model using Hugging Face pipeline
        self.model = pipeline("text2text-generation", model=self.model_name, device=self.device)

    def _get_binary_response(self, passage: str, query: str) -> str:
        """
        Performs binary classification to determine if a passage is relevant to the query.

        Parameters
        ----------
        passage : str
            The passage to evaluate.
        query : str
            The query against which relevance is judged.

        Returns
        -------
        str
            "yes" or "no" response indicating relevance.
        """
        prompt = f"Is the following passage related to the query?\npassage: {passage}\nquery: {query}\nAnswer in yes or no."
        return self.model(prompt)[0]["generated_text"].strip().lower()

    def _get_pairwise_response(self, query: str, passage_a: str, passage_b: str) -> str:
        """
        Compares two passages and determines which is more relevant to the query.

        Parameters
        ----------
        query : str
            The query for comparison.
        passage_a : str
            The first passage.
        passage_b : str
            The second passage.

        Returns
        -------
        str
            "passage a" or "passage b" indicating the more relevant passage.
        """
        prompt = f"""Given a query "{query}", which of the following two passages is more relevant to the query?
Passage A: {passage_a}
Passage B: {passage_b}
Output Passage A or Passage B."""
        return self.model(prompt)[0]["generated_text"].strip().lower()

    def rank(self, documents: List[Document]) -> List[Document]:
        """
        Reranks the contexts within each document using the EchoRank method.

        Parameters
        ----------
        documents : list of Document
            A list of documents containing query and contexts.

        Returns
        -------
        List[Document]
            The documents with reordered contexts based on their scores.

        Raises
        ------
        ValueError
            If no contexts are provided in a document.
        """
        for document in tqdm(documents, desc="Reranking Documents"):
            query = document.question.question
            contexts = document.contexts

            # Stage 1: Binary Classification
            binary_token_limit = int(self.budget_split_x * self.budget_tokens)
            binary_running_token = 0
            yes_contexts, no_contexts = [], []

            for context in contexts[:self.total_passages]:
                text = context.text
                token_length = len(text.split())
                if binary_running_token + token_length < binary_token_limit:
                    response = self._get_binary_response(text, query)
                    if "yes" in response:
                        yes_contexts.append(context)
                    else:
                        no_contexts.append(context)
                    binary_running_token += token_length

            # Stage 2: Pairwise Ranking
            pairwise_contexts = yes_contexts[:int(self.budget_split_y * self.budget_tokens)]
            for i in range(len(pairwise_contexts) - 1):
                for j in range(i + 1, len(pairwise_contexts)):
                    response = self._get_pairwise_response(query, pairwise_contexts[i].text, pairwise_contexts[j].text)
                    if response == "passage b":
                        pairwise_contexts[i], pairwise_contexts[j] = pairwise_contexts[j], pairwise_contexts[i]

            # Assign scores and sort
            all_contexts = pairwise_contexts + no_contexts
            for context in all_contexts:
                context.score = 1.0 if context in pairwise_contexts else 0.0  # Assigning scores based on pairwise ranking
            
            ranked_contexts = sorted(all_contexts, key=lambda ctx: ctx.score, reverse=True)

            # Update Document with reranked contexts
            document.reorder_contexts = ranked_contexts

        return documents

