from dataclasses import dataclass
from typing import List

from tvalmetrics.scorers.augmentation_accuracy_scorer import AugmentationAccuracyScorer
from tvalmetrics.scorers.retrieval_precision_scorer import RetrievalPrecisionScorer


@dataclass
class AugmentationPrecision:
    """Augmentation precision score and information used to calculate the score.

    Fields
    ------
    score: float
        Float between 0 and 1 representing the augmentation precision score.
    question: str
        The question that was asked.
    answer: str
        The answer that was generated by the RAG system.
    context_list: List[str]
        Retrieved context used by the RAG system to make answer.
    context_relevant_list: List[bool]
        List of booleans representing whether each context in context_list is relevant
        for answering the question.
    answer_contains_context_list: List[bool]
        List of booleans representing whether information from each context in
        context_list is in answer.
    """

    score: float
    question: str
    answer: str
    context_list: List[str]
    context_relevant_list: List[bool]
    answer_contains_context_list: List[bool]


class AugmentationPrecisionScorer(object):
    """Class for calculating augmentation precision score.

    Parameters
    ----------
    model: str
        Name of the LLM model to use as the LLM evaluator.
    """

    def __init__(self, model: str):
        self.model = model
        self.augmentation_accuracy_scorer = AugmentationAccuracyScorer(self.model)
        self.retrieval_precision_scorer = RetrievalPrecisionScorer(self.model)

    def score(
        self, question: str, answer: str, context_list: List[str]
    ) -> AugmentationPrecision:
        """Calculate augmentation precision score.

        Parameters
        ----------
        question: str
            The question that was asked.
        answer: str
            The answer that was generated by the RAG system.
        context_list: List[str]
            Retrieved context used by the RAG system to make answer.

        Returns
        -------
        AugmentationPrecisionScore
            Augmentation precision score and information used to calculate the score.
        """
        retrieval_precision_score = self.retrieval_precision_scorer.score(
            question, context_list
        )
        context_relevant_list = retrieval_precision_score.context_relevant_list
        augmentation_accuracy_score = self.augmentation_accuracy_scorer.score(
            answer, context_list
        )
        contains_context_list = augmentation_accuracy_score.answer_contains_context_list

        score = self.score_from_context_labels(
            context_relevant_list, contains_context_list
        )
        return AugmentationPrecision(
            score,
            answer,
            question,
            context_list,
            context_relevant_list,
            contains_context_list,
        )

    @staticmethod
    def score_from_context_labels(
        context_relevant_list: List[bool], answer_contains_context_list: List[bool]
    ) -> float:
        relevant_context_count = 0
        relevant_context_used_count = 0
        for context_relevant, contains_context in zip(
            context_relevant_list, answer_contains_context_list
        ):
            if context_relevant:
                relevant_context_count += 1
                if contains_context:
                    relevant_context_used_count += 1
        if relevant_context_count == 0:
            augmentation_precision = 0.0
        else:
            augmentation_precision = (
                relevant_context_used_count / relevant_context_count
            )

        return augmentation_precision
