from pathlib import Path
from typing import Any

import bert_score
from moonshot.src.metrics.metric_interface import MetricInterface
from moonshot.src.utils.timeit import timeit


class BertScore(MetricInterface):
    """
    BertScore uses Bert to check for the similarity in embedding between two sentences.
    Code reference from:
    https://github.com/Tiiiger/bert_score/blob/master/bert_score_cli/score.py
    """

    def __init__(self):
        self.id = Path(__file__).stem
        self.name = "BertScore"
        self.description = "BertScore uses Bert to check for the similarity in embedding between two sentences."
        self.metric_config = self.get_metrics_configuration(self.id)

    def get_metadata(self) -> dict | None:
        """
        Retrieves and returns the metadata of the BertScore class.
        The metadata includes the unique identifier, the name, and the description of the class.

        Returns:
            dict | None: A dictionary containing the 'id', 'name', and 'description' of the BertScore class,
            or None if not applicable.
        """
        return {"id": self.id, "name": self.name, "description": self.description}

    @timeit
    async def get_results(
        self, prompts: Any, predicted_results: Any, targets: Any, *args, **kwargs
    ) -> dict:
        """
        Calculate the BERTScore precision, recall, and F1 score between the predicted results and the target results.

        Parameters:
            prompts (Any): The prompts used for generating the predicted results.
            predicted_results (Any): The predicted results generated by the model.
            targets (Any): The target results for comparison.
            *args: Variable length argument list.
            **kwargs: Arbitrary keyword arguments.

        Returns:
            dict: A dictionary containing the BERTScore precision, recall, and F1 score.
                - precision (float): The precision score.
                - recall (float): The recall score.
                - f1 (float): The F1 score.
        """
        # use default roberto model
        score = bert_score.score(
            predicted_results, targets, lang="en", rescale_with_baseline=True
        )

        avg_scores = [s.mean(dim=0) for s in score]
        precision_value = avg_scores[0].cpu().item()
        recall_value = avg_scores[1].cpu().item()
        f1_value = avg_scores[2].cpu().item()

        return {
            "bertscore": {
                "precision": precision_value,
                "recall": recall_value,
                "f1": f1_value,
            },
            "grading_criteria": {},
        }
