"""
Adapted from https://github.com/AIR-Bench/AIR-Bench/blob/0.1.0/air_benchmark/evaluation_utils/evaluator.py
"""
import json
import logging
import os
import json
import pandas as pd
from typing import Dict, Optional, List, Union

from .data_loader import AbsEvalDataLoader
from .searcher import EvalRetriever, EvalReranker
from .utils import evaluate_metrics, evaluate_mrr

logger = logging.getLogger(__name__)


class AbsEvaluator:
    """
    Base class of Evaluator.
    
    Args:
        eval_name (str): The experiment name of current evaluation.
        data_loader (AbsEvalDataLoader): The data_loader to deal with data.
        overwrite (bool): If true, will overwrite the existing results.
    """
    def __init__(
        self,
        eval_name: str,
        data_loader: AbsEvalDataLoader,
        overwrite: bool = False,
    ):
        self.eval_name = eval_name
        self.data_loader = data_loader
        self.overwrite = overwrite

    def check_data_info(
        self,
        data_info: Dict[str, str],
        model_name: str,
        reranker_name: str,
        split: str,
        dataset_name: Optional[str] = None,
    ):
        """Check the validity of data info.

        Args:
            data_info (Dict[str, str]): The loaded data info to be check.
            model_name (str): Name of model used.
            reranker_name (str): Name of reranker used.
            split (str): Split used in searching.
            dataset_name (Optional[str], optional): Name of dataset used. Defaults to None.

        Raises:
            ValueError: eval_name mismatch
            ValueError: model_name or reranker_name mismatch
            ValueError: split mismatch
            ValueError: dataset_name mismatch
        """
        if data_info["eval_name"] != self.eval_name:
            raise ValueError(
                f'eval_name mismatch: {data_info["eval_name"]} vs {self.eval_name}'
            )
        if (
            data_info["model_name"] != model_name
            or data_info["reranker_name"] != reranker_name
        ):
            raise ValueError(
                f'model_name or reranker_name mismatch: {data_info["model_name"]} vs {model_name} or {data_info["reranker_name"]} vs {reranker_name}'
            )
        if (data_info["split"] != split):
            raise ValueError(
                f'split mismatch: {data_info["split"]} vs {split}'
            )
        if dataset_name is not None and data_info["dataset_name"] != dataset_name:
            raise ValueError(
                f'dataset_name mismatch: {data_info["dataset_name"]} vs {dataset_name}'
            )

    def get_corpus_embd_save_dir(
        self,
        retriever_name: str,
        corpus_embd_save_dir: Optional[str] = None,
        dataset_name: Optional[str] = None
    ):
        """
        If corpus_embd_save_dir is not None, then it will be used as the base directory to save the corpus embeddings. For dataset such as MKQA, 
            the corpus for all languages is the same, so the subclass can override this method to save the corpus embeddings in the same directory.
        
        Args:
            retriever_name (str): Name of the retriever.
            corpus_embd_save_dir (str, optional): Directory that saving the corpus embedding.
            dataset_name (str, optional): 
        """
        if corpus_embd_save_dir is not None:
            if dataset_name is not None:
                corpus_embd_save_dir = os.path.join(corpus_embd_save_dir, retriever_name, dataset_name)
            else:
                corpus_embd_save_dir = os.path.join(corpus_embd_save_dir, retriever_name)
        return corpus_embd_save_dir

    def __call__(
        self,
        splits: Union[str, List[str]],
        search_results_save_dir: str,
        retriever: EvalRetriever,
        reranker: Optional[EvalReranker] = None,
        corpus_embd_save_dir: Optional[str] = None,
        ignore_identical_ids: bool = False,
        k_values: List[int] = [1, 3, 5, 10, 100, 1000],
        dataset_name: Optional[str] = None,
        **kwargs,
    ):
        """This is called during the evaluation process.

        Args:
            splits (Union[str, List[str]]): Splits of datasets.
            search_results_save_dir (str): Directory to save the search results.
            retriever (EvalRetriever): object of :class:EvalRetriever.
            reranker (Optional[EvalReranker], optional): Object of :class:EvalReranker. Defaults to :data:`None`.
            corpus_embd_save_dir (Optional[str], optional): Directory to save the embedded corpus. Defaults to :data:`None`.
            ignore_identical_ids (bool, optional): If True, will ignore identical ids in search results. Defaults to :data:`False`.
            k_values (List[int], optional): Cutoffs. Defaults to :data:`[1, 3, 5, 10, 100, 1000]`.
            dataset_name (Optional[str], optional): Name of the datasets. Defaults to :data:`None`.
        """
        # Check Splits
        checked_splits = self.data_loader.check_splits(splits, dataset_name=dataset_name)
        if len(checked_splits) == 0:
            logger.warning(f"{splits} not found in the dataset. Skipping evaluation.")
            return
        splits = checked_splits

        if dataset_name is not None:
            save_name = f"{dataset_name}-" + "{split}.json"
        else:
            save_name = "{split}.json"

        corpus_embd_save_dir = self.get_corpus_embd_save_dir(
            retriever_name=str(retriever),
            corpus_embd_save_dir=corpus_embd_save_dir,
            dataset_name=dataset_name
        )

        # Retrieval Stage
        no_reranker_search_results_save_dir = os.path.join(
            search_results_save_dir, str(retriever), "NoReranker"
        )
        os.makedirs(no_reranker_search_results_save_dir, exist_ok=True)

        flag = False
        for split in splits:
            split_no_reranker_search_results_save_path = os.path.join(
                no_reranker_search_results_save_dir, save_name.format(split=split)
            )
            if not os.path.exists(split_no_reranker_search_results_save_path) or self.overwrite:
                flag = True
                break

        no_reranker_search_results_dict = {}
        if flag:
            corpus = self.data_loader.load_corpus(dataset_name=dataset_name)

            queries_dict = {
                split: self.data_loader.load_queries(dataset_name=dataset_name, split=split)
                for split in splits
            }

            all_queries = {}
            for _, split_queries in queries_dict.items():
                all_queries.update(split_queries)

            all_no_reranker_search_results = retriever(
                corpus=corpus,
                queries=all_queries,
                corpus_embd_save_dir=corpus_embd_save_dir,
                ignore_identical_ids=ignore_identical_ids,
                **kwargs,
            )

            for split in splits:
                split_queries = queries_dict[split]
                no_reranker_search_results_dict[split] = {
                    qid: all_no_reranker_search_results[qid] for qid in split_queries
                }
                split_no_reranker_search_results_save_path = os.path.join(
                    no_reranker_search_results_save_dir, save_name.format(split=split)
                )

                self.save_search_results(
                    eval_name=self.eval_name,
                    model_name=str(retriever),
                    reranker_name="NoReranker",
                    search_results=no_reranker_search_results_dict[split],
                    output_path=split_no_reranker_search_results_save_path,
                    split=split,
                    dataset_name=dataset_name,
                )
        else:
            for split in splits:
                split_no_reranker_search_results_save_path = os.path.join(
                    no_reranker_search_results_save_dir, save_name.format(split=split)
                )
                data_info, search_results = self.load_search_results(split_no_reranker_search_results_save_path)

                self.check_data_info(
                    data_info=data_info,
                    model_name=str(retriever),
                    reranker_name="NoReranker",
                    split=split,
                    dataset_name=dataset_name,
                )
                no_reranker_search_results_dict[split] = search_results
        retriever.stop_multi_process_pool()
        eval_results_save_path = os.path.join(no_reranker_search_results_save_dir, 'EVAL', 'eval_results.json')
        if not os.path.exists(eval_results_save_path) or self.overwrite or flag:
            retriever_eval_results = self.evaluate_results(no_reranker_search_results_save_dir, k_values=k_values)
            self.output_eval_results_to_json(retriever_eval_results, eval_results_save_path)

        # Reranking Stage
        if reranker is not None:
            reranker_search_results_save_dir = os.path.join(
                search_results_save_dir, str(retriever), str(reranker)
            )
            os.makedirs(reranker_search_results_save_dir, exist_ok=True)

            corpus = self.data_loader.load_corpus(dataset_name=dataset_name)

            queries_dict = {
                split: self.data_loader.load_queries(dataset_name=dataset_name, split=split)
                for split in splits
            }

            flag = False
            for split in splits:
                rerank_search_results_save_path = os.path.join(
                    reranker_search_results_save_dir, save_name.format(split=split)
                )

                if os.path.exists(rerank_search_results_save_path) and not self.overwrite:
                    continue

                flag = True
                rerank_search_results = reranker(
                    corpus=corpus,
                    queries=queries_dict[split],
                    search_results=no_reranker_search_results_dict[split],
                    ignore_identical_ids=ignore_identical_ids,
                    **kwargs,
                )

                self.save_search_results(
                    eval_name=self.eval_name,
                    model_name=str(retriever),
                    reranker_name=str(reranker),
                    search_results=rerank_search_results,
                    output_path=rerank_search_results_save_path,
                    split=split,
                    dataset_name=dataset_name,
                )
            reranker.stop_multi_process_pool()
            eval_results_save_path = os.path.join(reranker_search_results_save_dir, 'EVAL', 'eval_results.json')
            if not os.path.exists(eval_results_save_path) or self.overwrite or flag:
                reranker_eval_results = self.evaluate_results(reranker_search_results_save_dir, k_values=k_values)
                self.output_eval_results_to_json(reranker_eval_results, eval_results_save_path)

    @staticmethod
    def save_search_results(
        eval_name: str,
        model_name: str,
        reranker_name: str,
        search_results: Dict[str, Dict[str, float]],
        output_path: str,
        split: str,
        dataset_name: Optional[str] = None,
    ):
        """Save the metadata and search results into a file.

        Args:
            eval_name (str): The experiment name of current evaluation.
            model_name (str): Name of model used.
            reranker_name (str): Name of reranker used.
            search_results (Dict[str, Dict[str, float]]): Dictionary of search results.
            output_path (str): Output path to write the results.
            split (str): Split used in searching.
            dataset_name (Optional[str], optional): Name of dataset used. Defaults to :data:`None`.
        """
        data = {
            "eval_name": eval_name,
            "model_name": model_name,
            "reranker_name": reranker_name,
            "split": split,
            "dataset_name": dataset_name,
            "search_results": search_results,
        }

        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=4)

    @staticmethod
    def load_search_results(input_path: str):
        """Load search results from path.

        Args:
            input_path (str): Path to load from.

        Returns:
            dict, dict: data info that contains metadata and search results.
        """
        with open(input_path, "r", encoding="utf-8") as f:
            data_info = json.load(f)
        
        search_results = data_info.pop("search_results")
        return data_info, search_results

    @staticmethod
    def compute_metrics(
        qrels: Dict[str, Dict[str, int]],
        search_results: Dict[str, Dict[str, float]],
        k_values: List[int],
    ):
        """Evaluate the model with metrics.

        Args:
            qrels (Dict[str, Dict[str, int]]): Ground truth relevance of queries and documents.
            search_results (Dict[str, Dict[str, float]]): Dictionary of search results
            k_values (List[int]): Cutoffs.

        Returns:
            dict: The results of the metrics.
        """
        ndcg, _map, recall, precision = evaluate_metrics(
            qrels=qrels,
            results=search_results,
            k_values=k_values,
        )
        mrr = evaluate_mrr(
            qrels=qrels,
            results=search_results,
            k_values=k_values,
        )
        scores = {
            **{f"ndcg_at_{k.split('@')[1]}": v for (k, v) in ndcg.items()},
            **{f"map_at_{k.split('@')[1]}": v for (k, v) in _map.items()},
            **{f"recall_at_{k.split('@')[1]}": v for (k, v) in recall.items()},
            **{f"precision_at_{k.split('@')[1]}": v for (k, v) in precision.items()},
            **{f"mrr_at_{k.split('@')[1]}": v for (k, v) in mrr.items()},
        }
        return scores

    def evaluate_results(
        self,
        search_results_save_dir: str,
        k_values: List[int] = [1, 3, 5, 10, 100, 1000]
    ):
        """Compute metrics according to the results in the directory.

        Args:
            search_results_save_dir (str): Path to the search results.
            k_values (List[int], optional): Cutoffs. Defaults to :data:`[1, 3, 5, 10, 100, 1000]`.

        Returns:
            dict: Evaluation results.
        """
        eval_results_dict = {}

        for file in os.listdir(search_results_save_dir):
            if not file.endswith('.json'):
                continue

            file_path = os.path.join(search_results_save_dir, file)
            data_info, search_results = self.load_search_results(file_path)

            _eval_name = data_info['eval_name']
            assert _eval_name == self.eval_name, f'Mismatch eval_name: {_eval_name} vs {self.eval_name} in {file_path}'

            split = data_info['split']
            dataset_name = data_info.get('dataset_name', None)
            qrels = self.data_loader.load_qrels(dataset_name=dataset_name, split=split)

            eval_results = self.compute_metrics(
                qrels=qrels,
                search_results=search_results,
                k_values=k_values
            )

            if dataset_name is not None:
                key = f"{dataset_name}-{split}"
            else:
                key = split
            eval_results_dict[key] = eval_results

        return eval_results_dict

    @staticmethod
    def output_eval_results_to_json(eval_results_dict: dict, output_path: str):
        """Write the evaluation results into a json file.

        Args:
            eval_results_dict (dict): Dictionary of the evaluation results.
            output_path (str): Output path to write the json file.
        """
        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(eval_results_dict, f, indent=4)
        logger.info(f"Results saved to {output_path}")

    @staticmethod
    def get_results_df(metric: str, eval_results_dict: dict):
        """Get the results from dictionary to a DataFrame.

        Args:
            metric (str): Selected metric.
            eval_results_dict (dict): Dictionary of the evaluation results.

        Returns:
            DataFrame: DataFrame of the results.
        """
        results_dict = {}

        for model_name, model_results in eval_results_dict.items():
            results_dict[model_name] = {}
            for reranker_name, reranker_results in model_results.items():
                results_dict[model_name][reranker_name] = {}
                for split, split_results in reranker_results.items():
                    if metric in split_results:
                        results_dict[model_name][reranker_name][split] = split_results[metric]
                    else:
                        results_dict[model_name][reranker_name][split] = None

        model_reranker_pairs = set()
        all_splits = set()
        for model_name, model_results in results_dict.items():
            for reranker_name, reranker_results in model_results.items():
                model_reranker_pairs.add((model_name, reranker_name))
                all_splits.update(reranker_results.keys())

        index = [(model, reranker) for model, reranker in model_reranker_pairs]
        multi_index = pd.MultiIndex.from_tuples(index, names=['Model', 'Reranker'])
        
        all_splits = sorted(list(all_splits))
        overall_columns = ['average'] + all_splits
        overall_df = pd.DataFrame(index=multi_index, columns=overall_columns)
        
        for model, reranker in model_reranker_pairs:
            for split in all_splits:
                if model in results_dict and reranker in results_dict[model] and split in results_dict[model][reranker]:
                    overall_df.loc[(model, reranker), split] = results_dict[model][reranker][split]
                else:
                    overall_df.loc[(model, reranker), split] = None
            if overall_df.loc[(model, reranker), all_splits].isnull().any():
                overall_df.loc[(model, reranker), 'average'] = None
            else:
                overall_df.loc[(model, reranker), 'average'] = overall_df.loc[(model, reranker), all_splits].mean()

        return overall_df

    @staticmethod
    def output_eval_results_to_markdown(eval_results_dict: dict, output_path: str, metrics: Union[List[str], str]):
        """Write the evaluation results to a markdown file.

        Args:
            eval_results_dict (dict): Dictionary that contains evaluation results.
            output_path (str): Path to write the output to.
            metrics (Union[List[str], str]): The metrics that will be written in the markdown file.
        """
        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        if isinstance(metrics, str):
            metrics = [metrics]

        with open(output_path, 'w', encoding='utf-8') as f:
            for metric in metrics:
                f.write(f"## {metric}\n\n")
                results_df = AbsEvaluator.get_results_df(metric, eval_results_dict)
                max_index = dict(results_df.idxmax(axis=0))
                splits = results_df.columns
                f.write(f"| Model | Reranker | {' | '.join(splits)} |\n")
                f.write(f"| :---- | :---- | {' | '.join([':---:' for _ in splits])} |\n")
                for i, row in results_df.iterrows():
                    line = f"| {i[0]} | {i[1]} | "
                    for s, v in row.items():
                        if v is None:
                            line += "- | "
                        else:
                            if i != max_index[s]:
                                line += f'{v*100:.3f} | '
                            else:
                                line += f'**{v*100:.3f}** | '
                    f.write(line + "\n")
                f.write("\n")
        logger.info(f"Results saved to {output_path}")
