"""
Author: Ronen Huang
"""

import os
import multiprocessing
import lightgbm
import pandas as pd
from tqdm import tqdm
import numpy as np
import optuna
import scipy.stats as st
from bible_search_engine.components.indexing import BibleChapterIndex
from bible_search_engine.components.preprocess import NLPTokenizer
from bible_search_engine.components.ranker import\
    TraditionalRanker, TFIDFRanker, BM25Ranker, \
    DirichletLMRanker, BiEncoderRanker, ColbertRanker
from bible_search_engine.components.relevance import Relevance


class L2RRanker:
    """
    Learning to Rank Bible chapter retrieval.
    """
    def __init__(self, bible_chapter_index: BibleChapterIndex,
                 nlp_tokenizer: NLPTokenizer,
                 ranker: TraditionalRanker | ColbertRanker,
                 feature_extractor: 'L2RFeatureExtractor') -> None:
        '''
        bible_chapter_index: Bible chapter index.
        nlp_tokenizer: NLP tokenizer.
        ranker: Ranker with query.
        feature_extractor: Feature extractor for LightGBM ranker.

        Initializes learning to rank with LightGBM ranker.
        '''
        self.bible_chapter_index = bible_chapter_index
        self.nlp_tokenizer = nlp_tokenizer
        self.ranker = ranker
        self.feature_extractor = feature_extractor
        self.lightgbm_ranker = None

    def train(self, train_queries_data: str, test_queries_data: str) -> None:
        '''
        Train LightGBM ranker with train queries and their relevance scores.
        '''
        if not os.path.isfile(train_queries_data):
            raise Exception('Train queries data does not exist.')
        if not os.path.isfile(test_queries_data):
            raise Exception('Test queries data does not exist.')

        train_eval = Relevance(train_queries_data)

        test_eval = Relevance(test_queries_data)

        params_study = optuna.create_study(
            study_name="Bible Search Engine Ranker Parameter Optimization",
            direction="maximize"
        )
        params_study.optimize(
            lambda trial: self.determine_params(
                trial, train_eval, test_eval
            ), 20
        )
        best_params = params_study.best_params

        lgbmranker_params = {
            'num_leaves': best_params["num_leaves"],
            'learning_rate': best_params["learning_rate"],
            'max_depth': best_params["max_depth"],
            'num_iterations': best_params["num_iterations"],
            'min_data_in_leaf': best_params["min_data_in_leaf"],
            'metric': 'ndcg', 'importance_type': 'gain', 'verbosity': -1,
            'n_jobs': multiprocessing.cpu_count()}
        self.lightgbm_ranker =\
            lightgbm.LGBMRanker(**lgbmranker_params)
        self.feature_extractor.bm25_ranker.set_params(
            best_params["b"], best_params["k1"], best_params["k3"]
        )
        self.feature_extractor.dirichlet_lm_ranker.set_params(
            best_params["mu"]
        )

        queries_df = pd.concat(
            [train_eval.relevance_df, test_eval.relevance_df],
            ignore_index=True
        )
        train_features, train_relevance_scores, train_num_query_examples =\
            self.get_train(queries_df)

        self.lightgbm_ranker.fit(
            train_features, train_relevance_scores,
            group=train_num_query_examples
        )

    def determine_params(self, trial: optuna.Trial,
                         train_eval: Relevance,
                         test_eval: Relevance) -> float:
        """
        Returns test precision for hyperparameters.
        """
        learning_rate = trial.suggest_float(
            "learning_rate", 5e-6, 5e-4, log=True
        )
        num_iterations = trial.suggest_int(
            "num_iterations", 1000, 2000, log=True
        )
        max_depth = trial.suggest_int("max_depth", 7, 14)
        num_leaves = trial.suggest_int(
            "num_leaves", 2 ** (max_depth - 2), 2 ** (max_depth - 1)
        )
        min_data_in_leaf = trial.suggest_int("min_data_in_leaf", 5, 35)
        b = trial.suggest_float("b", 0.3, 0.9)
        k1 = trial.suggest_float("k1", 0.8, 2.2)
        k3 = trial.suggest_float("k3", 0, 15)
        mu = trial.suggest_int("mu", 1000, 5000)

        self.feature_extractor.bm25_ranker.set_params(b, k1, k3)
        self.feature_extractor.dirichlet_lm_ranker.set_params(mu)

        train_features, train_relevance_scores, train_num_query_examples =\
            self.get_train(train_eval.relevance_df)

        lgbmranker_params = {
            'num_leaves': num_leaves, 'learning_rate': learning_rate,
            'max_depth': max_depth, 'num_iterations': num_iterations,
            'min_data_in_leaf': min_data_in_leaf, 'metric': 'ndcg',
            'importance_type': 'gain', 'verbosity': -1,
            'n_jobs': multiprocessing.cpu_count()
        }
        self.lightgbm_ranker =\
            lightgbm.LGBMRanker(**lgbmranker_params)
        self.lightgbm_ranker.fit(
            train_features, train_relevance_scores,
            group=train_num_query_examples
        )

        results = [s for _, s in test_eval.evaluate_ranker_results(self)]

        test_features, test_relevance_scores, test_num_query_examples =\
            self.get_train(test_eval.relevance_df)
        self.lightgbm_ranker =\
            lightgbm.LGBMRanker(**lgbmranker_params)
        self.lightgbm_ranker.fit(
            test_features, test_relevance_scores,
            group=test_num_query_examples
        )

        results.extend(
            [s for _, s in train_eval.evaluate_ranker_results(self)]
        )

        return st.t.interval(
            0.99, len(results) - 1, np.mean(results),
            st.sem(results)
        )[0]

    def get_train(self, train_queries_df: pd.DataFrame) ->\
            tuple[list[list[float]], list[int], list[int]]:
        """
        Prepares train data for L2R.
        """
        train_features = []
        train_relevance_scores = []
        train_num_query_examples = []

        for query in tqdm(train_queries_df['query'].unique()):
            train_query_df = train_queries_df[
                train_queries_df['query'] == query
            ][['chapterid', 'relevance']]
            train_num_query_examples.append(train_query_df.shape[0])
            query_parts = self.nlp_tokenizer.tokenize(query)
            set_query_parts = set(
                [
                    query_part
                    for query_part in query_parts if query_part
                ]
            )
            for chapterid, relevance_score in train_query_df.itertuples(
                index=False, name=None
            ):
                chapter_term_counts = {}
                for query_part in set_query_parts.intersection(
                    self.bible_chapter_index.get_chapter_vocab(chapterid)
                ):
                    chapter_term_counts[query_part] =\
                        self.bible_chapter_index.get_chapter_term_freq(
                            chapterid, query_part
                        )
                train_features.append(
                    self.feature_extractor.get_features(
                        chapterid, chapter_term_counts, query_parts, query
                    )
                )
                train_relevance_scores.append(relevance_score)

        return train_features, train_relevance_scores, train_num_query_examples

    def query(self, query: str) -> list[tuple[int, float]]:
        '''
        query: Query of interest.

        Search for the relevant Bible chapters to the query.
        '''
        query_parts = self.nlp_tokenizer.tokenize(query)
        set_query_parts = set(
            [query_part for query_part in query_parts if query_part]
        )

        initial_results = self.ranker.query(query)
        if initial_results == []:
            return initial_results

        # Rerank top 200 initial results with learning to rank.
        top_200_initial_results = initial_results[:200]
        test_features = []
        for chapterid, _ in tqdm(top_200_initial_results):
            chapter_term_counts = {}
            for query_part in set_query_parts.intersection(
                self.bible_chapter_index.get_chapter_vocab(chapterid)
            ):
                chapter_term_counts[query_part] =\
                    self.bible_chapter_index.get_chapter_term_freq(
                        chapterid, query_part
                    )
            test_features.append(
                self.feature_extractor.get_features(
                    chapterid, chapter_term_counts, query_parts, query
                )
            )

        results = [
            top_200_initial_results[i]
            for i in np.argsort(
                self.lightgbm_ranker.predict(test_features)
            )[::-1]
        ]
        results.extend(initial_results[200:])

        return results

    def ranker_name(self) -> str:
        return 'l2r_ranker'


class L2RFeatureExtractor:
    """
    Learning to Rank feature extraction.
    """
    def __init__(self, bible_chapter_index: BibleChapterIndex,
                 nlp_tokenizer: NLPTokenizer, tf_idf_ranker: TFIDFRanker,
                 bm25_ranker: BM25Ranker,
                 dirichlet_lm_ranker: DirichletLMRanker,
                 biencoder_ranker: BiEncoderRanker) -> None:
        '''
        bible_chapter_index: Bible chapter index.
        nlp_tokenizer: NLP tokenizer.
        tf_idf_ranker: TF-IDF ranker with score.
        bm25_ranker: BM25 ranker with score.
        dirichlet_lm_ranker: Dirichlet LM ranker with score.
        '''
        self.bible_chapter_index = bible_chapter_index
        self.nlp_tokenizer = nlp_tokenizer
        self.tf_idf_ranker = tf_idf_ranker
        self.bm25_ranker = bm25_ranker
        self.dirichlet_lm_ranker = dirichlet_lm_ranker
        self.biencoder_ranker = biencoder_ranker
        # self.cross_encoder_ranker = cross_encoder_ranker
        # self.colbert_ranker = colbert_ranker

    def get_number_verses(self, chapterid: int) -> int:
        '''
        chapterid: Bible chapter id.

        Gets number of verses for the Bible chapter.
        '''
        bible_chapter_metadata =\
            self.bible_chapter_index.get_chapter_metadata(chapterid)
        if bible_chapter_metadata == {}:
            return 0
        return bible_chapter_metadata['number_of_verses']

    def get_chapter_length(self, chapterid: int) -> int:
        '''
        chapterid: Bible chapter id.

        Gets length of the Bible chapter.
        '''
        bible_chapter_metadata =\
            self.bible_chapter_index.get_chapter_metadata(chapterid)
        if bible_chapter_metadata == {}:
            return 0
        return bible_chapter_metadata['chapter_length']

    def get_tf_idf(self, chapterid: int, chapter_term_counts: dict[str, int],
                   query_parts: list[str]) -> float:
        '''
        chapterid: Bible chapter id.
        chapter_term_counts: Count of each term in Bible chapter.
        query_parts: Tokenized query.

        Gets TF-IDF score of the Bible chapter.
        '''
        return self.tf_idf_ranker.score(
            chapterid, chapter_term_counts, query_parts
        )

    def get_bm25(self, chapterid: int, chapter_term_counts: dict[str, int],
                 query_parts: list[str]) -> float:
        '''
        chapterid: Bible chapter id.
        chapter_term_counts: Count of each term in Bible chapter.
        query_parts: Tokenized query.

        Gets BM25 score of the Bible chapter.
        '''
        return self.bm25_ranker.score(
            chapterid, chapter_term_counts, query_parts
        )

    def get_dirichlet_lm(self, chapterid: int,
                         chapter_term_counts: dict[str, int],
                         query_parts: list[str]) -> float:
        '''
        chapterid: Bible chapter id.
        chapter_term_counts: Count of each term in Bible chapter.
        query_parts: Tokenized query.

        Gets Dirichlet LM score of the Bible chapter.
        '''
        return self.dirichlet_lm_ranker.score(
            chapterid, chapter_term_counts, query_parts
        )

    # def get_cross_encoder_score(self, chapterid: int, query: str) -> float:
    #    '''

    #    chapterid: Bible chapter id.
    #    query: Query of interest.

    #    Gets cross-encoder score of the Bible chapter.
    #    '''
    #    return self.cross_encoder_ranker.score(chapterid, query)

    # def get_colbert_score(self, chapterid: int, query: str) -> float:
    #     """
    #     chapterid: Bible chapter id.
    #     query: Query of interest.

    #     Gets colbert score of the Bible chapter.
    #     """
    #     return self.colbert_ranker.score(chapterid, query)

    def get_biencoder_score(self, chapterid: int, query: str) -> float:
        """
        chapterid: Bible chapter id.
        query: Query of interest.

        Gets biencoder score of the Bible chapter.
        """
        return self.biencoder_ranker.score(chapterid, query)

    def get_features(self, chapterid: int, chapter_term_counts: dict[str, int],
                     query_parts: list[str], query: str) -> list[float]:
        '''
        chapterid: Bible chapter id.
        chapter_term_counts: Count of each term in Bible chapter.
        query_parts: Tokenized query.
        query: Query of interest.

        Get features of the Bible chapter.
        '''
        bible_chapter_features = []

        # Number of verses.
        bible_chapter_features.append(self.get_number_verses(chapterid))

        # Bible chapter Length.
        bible_chapter_features.append(self.get_chapter_length(chapterid))

        # TF-IDF.
        bible_chapter_features.append(
            self.get_tf_idf(
                chapterid, chapter_term_counts, query_parts
            )
        )

        # BM25.
        bible_chapter_features.append(
            self.get_bm25(chapterid, chapter_term_counts, query_parts)
        )

        # Dirichlet LM.
        bible_chapter_features.append(
            self.get_dirichlet_lm(chapterid, chapter_term_counts, query_parts)
        )

        # Cross-encoder.
        # bible_chapter_features.append(
        #   self.get_cross_encoder_score(chapterid, query)
        # )

        # Biencoder.
        bible_chapter_features.append(
            self.get_biencoder_score(chapterid, query)
        )

        return bible_chapter_features
