from dataclasses import dataclass
from statistics import mean
from typing import Dict, List, Optional

import pandas as pd
from tvalmetrics.scorers import (  # type: ignore
    AnswerConsistencyBinaryScorer,
    AnswerConsistencyScorer,
    AnswerSimilarityScorer,
    ContextScorer,
    RetrievalKRecallScorer,
)


@dataclass
class Scores:
    """Scores that are calculated by RagScoresCalculator.batch_score.

    If an input was not used by RagScoresCalculator.score, then the corresponding
    attribute is a list of Nones.

    Fields
    ------
    scores_list: List[Scores]
        List of scores that were calculated for each question and answer.
    question_list: List[Optional[str]]
        List of questions in the batch.
    reference_answer_list: List[Optional[str]]
        List of reference answers in the batch.
    llm_answer_list: List[Optional[str]]
        List of answers generated by the RAG system for the questions in the batch.
    retrieved_context_list_list: List[Optional[List[str]]]
        List of retrieved contexts used by the RAG system to make answers to the
        questions in the batch.
    top_k_context_list_list: List[Optional[List[str]]]
        List of top k contexts that would be retrieved by the RAG system for the
        questions in the batch. There's an assumption that the retrieved context list is
        a subset of the top k context list.
    """

    answer_similarity_score_list: List[Optional[float]]
    retrieval_precision_list: List[Optional[float]]
    augmentation_precision_list: List[Optional[float]]
    augmentation_accuracy_list: List[Optional[float]]
    answer_consistency_list: List[Optional[float]]
    answer_consistency_binary_list: List[Optional[float]]
    retrieval_k_recall_list: List[Optional[float]]
    overall_score_list: List[Optional[float]]
    question_list: List[Optional[str]]
    reference_answer_list: List[Optional[str]]
    llm_answer_list: List[Optional[str]]
    retrieved_context_list_list: List[Optional[List[str]]]
    top_k_context_list_list: List[Optional[List[str]]]

    def to_dataframe(self) -> pd.DataFrame:
        """Convert Scores to a pandas DataFrame."""
        df = pd.DataFrame(
            {
                "question": [x for x in self.question_list],
                "reference_answer": [x for x in self.reference_answer_list],
                "llm_answer": [x for x in self.llm_answer_list],
                "retrieved_context": [x for x in self.retrieved_context_list_list],
                "top_k_context": [x for x in self.top_k_context_list_list],
                "answer_similarity_score": [
                    x for x in self.answer_similarity_score_list
                ],
                "retrieval_precision": [x for x in self.retrieval_precision_list],
                "augmentation_precision": [x for x in self.augmentation_precision_list],
                "augmentation_accuracy": [x for x in self.augmentation_accuracy_list],
                "answer_consistency": [x for x in self.answer_consistency_list],
                "answer_consistency_binary": [
                    x for x in self.answer_consistency_binary_list
                ],
                "retrieval_k_recall": [x for x in self.retrieval_k_recall_list],
                "overall_score": [x for x in self.overall_score_list],
            }
        )
        return df.dropna(axis=1, how="all")

    def mean_scores(self) -> Dict[str, Optional[float]]:
        """Calculate the mean of the non-None scores.

        Returns
        -------
        Dict[str, Union[float, None]]
            Dictionary mapping each score name to its mean score.
        """
        if self.answer_similarity_score_list[0] is not None:
            mean_answer_similarity_score = mean(  # type: ignore
                self.answer_similarity_score_list
            )
        else:
            mean_answer_similarity_score = None
        if self.retrieval_precision_list[0] is not None:
            mean_retrieval_precision = mean(  # type: ignore
                self.retrieval_precision_list
            )
        else:
            mean_retrieval_precision = None
        if self.augmentation_precision_list[0] is not None:
            mean_augmentation_precision = mean(  # type: ignore
                self.augmentation_precision_list
            )
        else:
            mean_augmentation_precision = None
        if self.augmentation_accuracy_list[0] is not None:
            mean_augmentation_accuracy = mean(  # type: ignore
                self.augmentation_accuracy_list
            )
        else:
            mean_augmentation_accuracy = None
        if self.answer_consistency_binary_list[0] is not None:
            mean_answer_consistency_binary = mean(  # type: ignore
                self.answer_consistency_binary_list
            )
        else:
            mean_answer_consistency_binary = None
        if self.answer_consistency_list[0] is not None:
            mean_answer_consistency = mean(self.answer_consistency_list)  # type: ignore
        else:
            mean_answer_consistency = None
        if self.retrieval_k_recall_list[0] is not None:
            mean_retrieval_k_recall = mean(self.retrieval_k_recall_list)  # type: ignore
        else:
            mean_retrieval_k_recall = None
        if self.overall_score_list[0] is not None:
            mean_overall_score = mean(self.overall_score_list)  # type: ignore
        else:
            mean_overall_score = None
        return dict(
            answer_similarity_score=mean_answer_similarity_score,
            retrieval_precision=mean_retrieval_precision,
            augmentation_precision=mean_augmentation_precision,
            augmentation_accuracy=mean_augmentation_accuracy,
            answer_consistency_binary=mean_answer_consistency_binary,
            answer_consistency=mean_answer_consistency,
            retrieval_k_recall=mean_retrieval_k_recall,
            overall_score=mean_overall_score,
        )

    def scores_to_dict(self) -> Dict[str, List[Optional[float]]]:
        return dict(
            answer_similarity_score=self.answer_similarity_score_list,
            retrieval_precision=self.retrieval_precision_list,
            augmentation_precision=self.augmentation_precision_list,
            augmentation_accuracy=self.augmentation_accuracy_list,
            answer_consistency_binary=self.answer_consistency_binary_list,
            answer_consistency=self.answer_consistency_list,
            retrieval_k_recall=self.retrieval_k_recall_list,
            overall_score=self.overall_score_list,
        )


class RagScoresCalculator(object):
    """Class for calculating any collection of RAG scores.

    Note: If all scores are set to False, then the default is to calculate the answer
    similarity score, retrieval precision, augmentation precision, augmentation
    accuracy, and answer consistency.

    Parameters
    ----------
    model: str
        Name of the LLM model to use as the LLM evaluator.
    answer_similarity_score: bool
        Whether to calculate the answer similarity score.
    retrieval_precision: bool
        Whether to calculate the retrieval precision score.
    augmentation_precision: bool
        Whether to calculate the augmentation precision score.
    augmentation_accuracy: bool
        Whether to calculate the augmentation accuracy score.
    answer_consistency: bool
        Whether to calculate the answer consistency score.
    answer_consistency_binary: bool
        Whether to calculate the answer consistency binary score.
    retrieval_k_recall: bool
        Whether to calculate the retrieval k-recall score.
    """

    def __init__(
        self,
        model: str,
        answer_similarity_score: bool = False,
        retrieval_precision: bool = False,
        augmentation_precision: bool = False,
        augmentation_accuracy: bool = False,
        answer_consistency: bool = False,
        answer_consistency_binary: bool = False,
        retrieval_k_recall: bool = False,
    ):
        # in this case we do the defualt because the user has not specified any scores
        if not (
            answer_similarity_score
            or retrieval_precision
            or augmentation_precision
            or augmentation_accuracy
            or answer_consistency
            or answer_consistency_binary
            or retrieval_k_recall
        ):
            answer_similarity_score = True
            retrieval_precision = True
            augmentation_precision = True
            augmentation_accuracy = True
            answer_consistency = True
        self.model = model
        self.answer_similarity_score = answer_similarity_score
        self.retrieval_precision = retrieval_precision
        self.augmentation_precision = augmentation_precision
        self.augmentation_accuracy = augmentation_accuracy
        self.answer_consistency_binary = answer_consistency_binary
        self.answer_consistency = answer_consistency
        self.retrieval_k_recall = retrieval_k_recall

        if answer_similarity_score:
            self.answer_similarity_scorer = AnswerSimilarityScorer(self.model)
        else:
            self.answer_similarity_scorer = None  # type: ignore

        if retrieval_precision and augmentation_precision and augmentation_accuracy:
            self.context_scorer = ContextScorer(self.model)
        elif not (
            retrieval_precision and augmentation_precision and augmentation_accuracy
        ):
            self.context_scorer = None  # type: ignore
        else:
            error_message = (
                "must have all three of retrieval_precision, augmentation_precision, "
                "and augmentation_accuracy set to True or False"
            )
            raise ValueError(error_message)

        if answer_consistency_binary:
            self.answer_consistency_binary_scorer = AnswerConsistencyBinaryScorer(
                self.model
            )
        else:
            self.answer_consistency_binary_scorer = None  # type: ignore

        if answer_consistency:
            self.answer_consistency_scorer = AnswerConsistencyScorer(self.model)
        else:
            self.answer_consistency_scorer = None  # type: ignore
        if retrieval_k_recall:
            self.retrieval_k_recall_scorer = RetrievalKRecallScorer(self.model)
        else:
            self.retrieval_k_recall_scorer = None  # type: ignore

    def validate_score_parameters(
        self,
        question: Optional[str] = None,
        reference_answer: Optional[str] = None,
        llm_answer: Optional[str] = None,
        retrieved_context_list: Optional[List[str]] = None,
        top_k_context_list: Optional[List[str]] = None,
    ) -> None:
        if self.answer_similarity_scorer is not None:
            if question is None or reference_answer is None or llm_answer is None:
                error_message = (
                    "Need question, reference_answer, and llm_answer for similarity "
                    "scorer"
                )
                raise ValueError(error_message)
        if self.context_scorer is not None:
            if question is None or llm_answer is None or retrieved_context_list is None:
                error_message = (
                    "Need question, llm_answer, and retrieved_context_list for "
                    "context scores scorer"
                )
                raise ValueError(error_message)
        if self.answer_consistency_binary_scorer is not None:
            if llm_answer is None or retrieved_context_list is None:
                error_message = (
                    "Need llm_answer and retrieved_context_list for answer consistency "
                    "binary scorer"
                )
                raise ValueError(error_message)
        if self.answer_consistency_scorer is not None:
            if llm_answer is None or retrieved_context_list is None:
                error_message = (
                    "Need llm_answer and retrieved_context_list for answer consistency "
                    "scorer"
                )
                raise ValueError(error_message)
        if self.retrieval_k_recall_scorer is not None:
            if (
                question is None
                or retrieved_context_list is None
                or top_k_context_list is None
            ):
                error_message = (
                    "Need question, retrieved_context_list, and top_k_context_list for "
                    "retrieval k recall scorer"
                )
                raise ValueError(error_message)

    def calculate_overall_score(
        self,
        answer_similarity_score: Optional[float],
        retrieval_precision: Optional[float],
        augmentation_precision: Optional[float],
        augmentation_accuracy: Optional[float],
        answer_consistency_binary: Optional[float],
        answer_consistency: Optional[float],
        retrieval_k_recall: Optional[float],
    ) -> float:
        score_sum = 0.0
        score_count = 0
        if answer_similarity_score is not None:
            score_sum += answer_similarity_score / 5
            score_count += 1
        if retrieval_precision is not None:
            score_sum += retrieval_precision
            score_count += 1
        if augmentation_precision is not None:
            score_sum += augmentation_precision
            score_count += 1
        if augmentation_accuracy is not None:
            score_sum += augmentation_accuracy
            score_count += 1
        if answer_consistency_binary is not None:
            score_sum += answer_consistency_binary
            score_count += 1
        if answer_consistency is not None:
            score_sum += answer_consistency
            score_count += 1
        if retrieval_k_recall is not None:
            score_sum += retrieval_k_recall
            score_count += 1
        return score_sum / score_count

    def score(
        self,
        question: Optional[str] = None,
        reference_answer: Optional[str] = None,
        llm_answer: Optional[str] = None,
        retrieved_context_list: Optional[List[str]] = None,
        top_k_context_list: Optional[List[str]] = None,
    ) -> Scores:
        """Calculate RAG scores.

        Note: Only the inputs needed to calculate the scores set in __init__ are
        necessary.

        Parameters
        ----------
        question: Optional[str]
            The question that was asked.
        reference_answer: Optional[str]
            The answer that was expected.
        llm_answer: Optional[str]
            The answer that was generated by the RAG system.
        retrieved_context_list: Optional[List[str]]
            Retrieved context used by the RAG system to make answer.
        top_k_context_list: Optional[List[str]]
            Top k contexts that would be retrieved by the RAG system. There's an
            assumption that the retrieved context list is a subset of the top k context
            list.

        Returns
        -------
        Scores
            All the scores that were calculated.
        """
        self.validate_score_parameters(
            question,
            reference_answer,
            llm_answer,
            retrieved_context_list,
            top_k_context_list,
        )
        if self.answer_similarity_scorer is not None:
            answer_similarity_score = self.answer_similarity_scorer.score(
                question, reference_answer, llm_answer  # type: ignore
            )
        else:
            answer_similarity_score = None

        if self.context_scorer is not None:
            context_scores = self.context_scorer.score(
                question, llm_answer, retrieved_context_list  # type: ignore
            )
            retrieval_precision = context_scores.retrieval_precision
            augmentation_precision = context_scores.augmentation_precision
            augmentation_accuracy = context_scores.augmentation_accuracy
        else:
            retrieval_precision = None
            augmentation_precision = None
            augmentation_accuracy = None

        if self.answer_consistency_binary_scorer is not None:
            answer_consistency_binary = self.answer_consistency_binary_scorer.score(
                llm_answer, retrieved_context_list  # type: ignore
            )
        else:
            answer_consistency_binary = None

        if self.answer_consistency_scorer is not None:
            answer_consistency_score = self.answer_consistency_scorer.score(
                llm_answer, retrieved_context_list  # type: ignore
            )
            annswer_consistency = answer_consistency_score.score
        else:
            annswer_consistency = None

        if self.retrieval_k_recall_scorer is not None:
            if top_k_context_list is None:
                raise ValueError(
                    "top_k_context_list must be provided if retrieval_k_recall is True"
                )
            retrieval_k_recall_score = self.retrieval_k_recall_scorer.score(
                question, retrieved_context_list, top_k_context_list  # type: ignore
            )
            retrieval_k_recall = retrieval_k_recall_score.score
        else:
            retrieval_k_recall = None

        overall_score = self.calculate_overall_score(
            answer_similarity_score,
            retrieval_precision,
            augmentation_precision,
            augmentation_accuracy,
            answer_consistency_binary,
            annswer_consistency,
            retrieval_k_recall,
        )

        return Scores(
            answer_similarity_score_list=[answer_similarity_score],
            retrieval_precision_list=[retrieval_precision],
            augmentation_precision_list=[augmentation_precision],
            augmentation_accuracy_list=[augmentation_accuracy],
            answer_consistency_binary_list=[answer_consistency_binary],
            answer_consistency_list=[annswer_consistency],
            retrieval_k_recall_list=[retrieval_k_recall],
            overall_score_list=[overall_score],
            question_list=[question],
            reference_answer_list=[reference_answer],
            llm_answer_list=[llm_answer],
            retrieved_context_list_list=[retrieved_context_list],
            top_k_context_list_list=[top_k_context_list],
        )

    def get_n_samples(
        self,
        question_list: Optional[List[Optional[str]]],
        retrieved_context_list_list: Optional[List[Optional[List[str]]]],
    ) -> int:
        if question_list is not None:
            n_samples = len(question_list)
        elif retrieved_context_list_list is not None:
            n_samples = len(retrieved_context_list_list)
        else:
            error_message = (
                "Need either question_list or retrieved_context_list_list to calculate "
                "metrics. Got neither."
            )
            raise ValueError(error_message)
        return n_samples

    def score_batch(
        self,
        question_list: Optional[List[Optional[str]]] = None,
        reference_answer_list: Optional[List[Optional[str]]] = None,
        llm_answer_list: Optional[List[Optional[str]]] = None,
        retrieved_context_list_list: Optional[List[Optional[List[str]]]] = None,
        top_k_context_list_list: Optional[List[Optional[List[str]]]] = None,
    ) -> Scores:
        """Calculate RAG scores for a batch of questions and answers.

        Notes:
            - Only the inputs needed to calculate the scores set in __init__ are
              necessary.
            - All inputs are lists that are assumed to be the same size with
              corresponding elements in the same positions.

        Parameters
        ----------
        question_list: Optional[List[Optional[str]]]
            Batch of questions that were asked.
        reference_answer_list: Optional[List[Optional[str]]]
            Batch of answers that were expected.
        llm_answer_list: Optional[List[Optional[str]]]
            Batch of answers generated by the RAG system.
        retrieved_context_list_list: Optional[List[Optional[List[str]]]]
            Batch of retrieved contexts used by the RAG system to make answer.
        top_k_context_list_list: Optional[List[Optional[List[str]]]]
            Batch of top k contexts that would be retrieved by the RAG system. There's
            an assumption that the retrieved context list is a subset of the top k
            context list.

        Returns
        -------
        Scores
            All the scores that were calculated.
        """
        n_samples = self.get_n_samples(question_list, retrieved_context_list_list)
        if question_list is None:
            question_list = [None for _ in range(n_samples)]
        if reference_answer_list is None:
            reference_answer_list = [None for _ in range(n_samples)]
        if llm_answer_list is None:
            llm_answer_list = [None for _ in range(n_samples)]
        if retrieved_context_list_list is None:
            retrieved_context_list_list = [None for _ in range(n_samples)]
        if top_k_context_list_list is None:
            top_k_context_list_list = [None for _ in range(n_samples)]

        score_list = []
        for (
            question,
            reference_answer,
            llm_answer,
            retrieved_context_list,
            top_k_context_list,
        ) in zip(
            question_list,
            reference_answer_list,
            llm_answer_list,
            retrieved_context_list_list,
            top_k_context_list_list,
        ):
            score = self.score(
                question,
                reference_answer,
                llm_answer,
                retrieved_context_list,
                top_k_context_list,
            )
            score_list.append(score)

        return Scores(
            answer_similarity_score_list=[
                x.answer_similarity_score_list[0] for x in score_list
            ],
            retrieval_precision_list=[
                x.retrieval_precision_list[0] for x in score_list
            ],
            augmentation_precision_list=[
                x.augmentation_precision_list[0] for x in score_list
            ],
            augmentation_accuracy_list=[
                x.augmentation_accuracy_list[0] for x in score_list
            ],
            answer_consistency_binary_list=[
                x.answer_consistency_binary_list[0] for x in score_list
            ],
            answer_consistency_list=[x.answer_consistency_list[0] for x in score_list],
            retrieval_k_recall_list=[x.retrieval_k_recall_list[0] for x in score_list],
            overall_score_list=[x.overall_score_list[0] for x in score_list],
            question_list=question_list,
            reference_answer_list=reference_answer_list,
            llm_answer_list=llm_answer_list,
            retrieved_context_list_list=retrieved_context_list_list,
            top_k_context_list_list=top_k_context_list_list,
        )
