import pandas as pd
import prettytable
import time
from abc import ABC, abstractmethod
from collections import OrderedDict
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.base import BaseEstimator
from .vectorizers import BM25Vectorizer


class BaseRetriever(BaseEstimator, ABC):
    """
    Abstract base class for all Retriever classes.
    All retrievers should inherit from this class.
    Each retriever class should implement a _fit_vectorizer method and a
    _compute_scores method
    """

    def __init__(self, vectorizer, top_n=10, verbose=False):
        self.vectorizer = vectorizer
        self.top_n = top_n
        self.verbose = verbose

    def fit(self, df: pd.DataFrame, y=None):
        """
        Fit the retriever to a list of documents or paragraphs

        Parameters
        ----------
        df: pandas.DataFrame object with all documents
        """
        self.metadata = df
        return self._fit_vectorizer(df)

    @abstractmethod
    def _fit_vectorizer(self, df):
        pass

    @abstractmethod
    def _compute_scores(self, query):
        pass

    def predict(self, query: str) -> OrderedDict:
        """
        Compute the top_n closest documents given a query

        Parameters
        ----------
        query: str

        Returns
        -------
        best_idx_scores: OrderedDict
            Dictionnaire with top_n best scores and idices of the documents as keys

        """
        t0 = time.time()
        scores = self._compute_scores(query)
        idx_scores = [(idx, score) for idx, score in enumerate(scores)]
        best_idx_scores = OrderedDict(
            sorted(idx_scores, key=(lambda tup: tup[1]), reverse=True)[: self.top_n]
        )

        # inspired from https://github.com/facebookresearch/DrQA/blob/50d0e49bb77fe0c6e881efb4b6fe2e61d3f92509/scripts/reader/interactive.py#L63
        if self.verbose:
            rank = 1
            table = prettytable.PrettyTable(["rank", "index", "title"])
            for i in range(len(closest_docs_indices)):
                index = closest_docs_indices[i]
                if self.paragraphs:
                    article_index = self.paragraphs[int(index)]["index"]
                    title = self.metadata.iloc[int(article_index)]["title"]
                else:
                    title = self.metadata.iloc[int(index)]["title"]
                table.add_row([rank, index, title])
                rank += 1
            print(table)
            print("Time: {} seconds".format(round(time.time() - t0, 5)))

        return best_idx_scores


class TfidfRetriever(BaseRetriever):
    """
    A scikit-learn estimator for TfidfRetriever. Trains a tf-idf matrix from a corpus
    of documents then finds the most N similar documents of a given input document by
    taking the dot product of the vectorized input document and the trained tf-idf matrix.

    Parameters
    ----------
    lowercase : boolean
        Convert all characters to lowercase before tokenizing. (default is True)
    preprocessor : callable or None
        Override the preprocessing (string transformation) stage while preserving
        the tokenizing and n-grams generation steps. (default is None)
    tokenizer : callable or None
        Override the string tokenization step while preserving the preprocessing
        and n-grams generation steps (default is None)
    stop_words : string {‘english’}, list, or None
        If a string, it is passed to _check_stop_list and the appropriate stop
        list is returned. ‘english’ is currently the only supported string value.
        If a list, that list is assumed to contain stop words, all of which will
        be removed from the resulting tokens.
        If None, no stop words will be used. max_df can be set to a value in the
        range [0.7, 1.0) to automatically detect and filter stop words based on
        intra corpus document frequency of terms.
        (default is None)
    token_pattern : string
        Regular expression denoting what constitutes a “token”. The default regexp
        selects tokens of 2 or more alphanumeric characters (punctuation is completely
        ignored and always treated as a token separator).
    ngram_range : tuple (min_n, max_n)
        The lower and upper boundary of the range of n-values for different n-grams
        to be extracted. All values of n such that min_n <= n <= max_n will be used.
        (default is (1, 1))
    max_df : float in range [0.0, 1.0] or int
        When building the vocabulary ignore terms that have a document frequency strictly
        higher than the given threshold (corpus-specific stop words). If float, the parameter
        represents a proportion of documents, integer absolute counts. This parameter is
        ignored if vocabulary is not None. (default is 1.0)
    min_df : float in range [0.0, 1.0] or int
        When building the vocabulary ignore terms that have a document frequency
        strictly lower than the given threshold. This value is also called cut-off
        in the literature. If float, the parameter represents a proportion of
        documents, integer absolute counts. This parameter is ignored if vocabulary
        is not None. (default is 1)
    vocabulary : Mapping or iterable, optional
        Either a Mapping (e.g., a dict) where keys are terms and values are indices
        in the feature matrix, or an iterable over terms. If not given, a vocabulary
        is determined from the input documents. (default is None)
    paragraphs : iterable
        an iterable which yields either str, unicode or file objects
    top_n : int (default 20)
        maximum number of top articles (or paragraphs) to retrieve
    verbose : bool, optional
        If true, all of the warnings related to data processing will be printed.

    Attributes
    ----------
    vectorizer : TfidfVectorizer
        See https://scikit-learn.org/stable/modules/generated/sklearn.feature_extraction.text.TfidfVectorizer.html
    tfidf_matrix : sparse matrix, [n_samples, n_features]
        Tf-idf-weighted document-term matrix.

    Examples
    --------
    >>> from pvqa.retriever import TfidfRetriever

    >>> retriever = TfidfRetriever(ngram_range=(1, 2), max_df=0.85, stop_words='english')
    >>> retriever.fit(X=df)
    >>> best_idx_scores = retriever.predict(X='How to grow cassava quickly?')
    """

    def __init__(
        self,
        lowercase=True,
        preprocessor=None,
        tokenizer=None,
        stop_words="english",
        token_pattern=r"(?u)\b\w\w+\b",
        ngram_range=(1, 2),
        max_df=0.85,
        min_df=2,
        vocabulary=None,
        top_n=20,
        verbose=False,
    ):
        self.lowercase = lowercase
        self.preprocessor = preprocessor
        self.tokenizer = tokenizer
        self.stop_words = stop_words
        self.token_pattern = token_pattern
        self.ngram_range = ngram_range
        self.max_df = max_df
        self.min_df = min_df
        self.vocabulary = vocabulary

        vectorizer = TfidfVectorizer(
            lowercase=self.lowercase,
            preprocessor=self.preprocessor,
            tokenizer=self.tokenizer,
            stop_words=self.stop_words,
            token_pattern=self.token_pattern,
            ngram_range=self.ngram_range,
            max_df=self.max_df,
            min_df=self.min_df,
            vocabulary=self.vocabulary,
        )
        super().__init__(vectorizer, top_n, verbose)

    def _fit_vectorizer(self, df, y=None):
        self.tfidf_matrix = self.vectorizer.fit_transform(df["content"])
        return self

    def _compute_scores(self, query):
        question_vector = self.vectorizer.transform([query])
        scores = self.tfidf_matrix.dot(question_vector.T).toarray()
        return scores


class BM25Retriever(BaseRetriever):
    """
    A scikit-learn estimator for BM25Retriever. Trains a matrix based on BM25 statistics
    from a corpus of documents then finds the most N similar documents of a given input
    query by computing the BM25 score for each document based on the query.

    Parameters
    ----------
    lowercase : boolean
        Convert all characters to lowercase before tokenizing. (default is True)
    preprocessor : callable or None
        Override the preprocessing (string transformation) stage while preserving
        the tokenizing and n-grams generation steps. (default is None)
    tokenizer : callable or None
        Override the string tokenization step while preserving the preprocessing
        and n-grams generation steps (default is None)
    stop_words : string {‘english’}, list, or None
        If a string, it is passed to _check_stop_list and the appropriate stop
        list is returned. ‘english’ is currently the only supported string value.
        If a list, that list is assumed to contain stop words, all of which will
        be removed from the resulting tokens.
        If None, no stop words will be used. max_df can be set to a value in the
        range [0.7, 1.0) to automatically detect and filter stop words based on
        intra corpus document frequency of terms.
        (default is None)
    token_pattern : string
        Regular expression denoting what constitutes a “token”. The default regexp
        selects tokens of 2 or more alphanumeric characters (punctuation is completely
        ignored and always treated as a token separator).
    ngram_range : tuple (min_n, max_n)
        The lower and upper boundary of the range of n-values for different n-grams
        to be extracted. All values of n such that min_n <= n <= max_n will be used.
        (default is (1, 1))
    max_df : float in range [0.0, 1.0] or int
        When building the vocabulary ignore terms that have a document frequency strictly
        higher than the given threshold (corpus-specific stop words). If float, the parameter
        represents a proportion of documents, integer absolute counts. This parameter is
        ignored if vocabulary is not None. (default is 1.0)
    min_df : float in range [0.0, 1.0] or int
        When building the vocabulary ignore terms that have a document frequency
        strictly lower than the given threshold. This value is also called cut-off
        in the literature. If float, the parameter represents a proportion of
        documents, integer absolute counts. This parameter is ignored if vocabulary
        is not None. (default is 1)
    vocabulary : Mapping or iterable, optional
        Either a Mapping (e.g., a dict) where keys are terms and values are indices
        in the feature matrix, or an iterable over terms. If not given, a vocabulary
        is determined from the input documents. (default is None)
    paragraphs : iterable
        an iterable which yields either str, unicode or file objects
    top_n : int (default 20)
        maximum number of top articles (or paragraphs) to retrieve
    verbose : bool, optional
        If true, all of the warnings related to data processing will be printed.
    k1 : float, optional (default=2.0)
        term k1 in the BM25 formula
    b : float, optional (default=0.75)
        term b in the BM25 formula
    floor : float or None, optional (default=None)
        floor value for idf terms

    Attributes
    ----------
    vectorizer : BM25Vectorizer

    Examples
    --------
    >>> from pvqa.retriever import BM25Retriever

    >>> retriever = BM25Retriever(ngram_range=(1, 2), max_df=0.85, stop_words='english')
    >>> retriever.fit(df=df)
    >>> best_idx_scores = retriever.predict(query='How to grow cassava quickly?')

    """

    def __init__(
        self,
        lowercase=True,
        preprocessor=None,
        tokenizer=None,
        stop_words="english",
        token_pattern=r"(?u)\b\w\w+\b",
        ngram_range=(1, 2),
        max_df=0.85,
        min_df=2,
        vocabulary=None,
        top_n=20,
        verbose=False,
        k1=2.0,
        b=0.75,
        floor=None,
    ):

        self.lowercase = lowercase
        self.preprocessor = preprocessor
        self.tokenizer = tokenizer
        self.stop_words = stop_words
        self.token_pattern = token_pattern
        self.ngram_range = ngram_range
        self.max_df = max_df
        self.min_df = min_df
        self.vocabulary = vocabulary
        self.k1 = k1
        self.b = b
        self.floor = floor

        vectorizer = BM25Vectorizer(
            lowercase=self.lowercase,
            preprocessor=self.preprocessor,
            tokenizer=self.tokenizer,
            stop_words=self.stop_words,
            token_pattern=self.token_pattern,
            ngram_range=self.ngram_range,
            max_df=self.max_df,
            min_df=self.min_df,
            vocabulary=self.vocabulary,
            k1=self.k1,
            b=self.b,
            floor=self.floor,
        )
        super().__init__(vectorizer, top_n, verbose)

    def _fit_vectorizer(self, df, y=None):
        self.bm25_matrix = self.vectorizer.fit_transform(df["content"])
        return self

    def _compute_scores(self, query):
        question_vector = self.vectorizer.transform([query], is_query=True)
        scores = self.bm25_matrix.dot(question_vector.T).toarray()
        return scores
