from pureml_llm.metric_base import MetricBase
from typing import Any
#from rouge_score import rouge_scorer
import nltk
from torchmetrics.text.rouge import ROUGEScore

class RougeScoreFmeasure(MetricBase):
    name: str = 'rouge_score_fmeasure'
    input_type: str = 'text'
    output_type: Any = None
    kwargs: Any = {}


    def parse_data(self,data):
        return data
    
    def compute(self,references,predictions,**kwargs):
        
        rouge_score = ROUGEScore()
        scores = rouge_score(preds = predictions, target = references)
        fmeasure = scores['rouge1_fmeasure']
        # recall = scores['rouge1_recall']
        # precision = scores['rouge1_precision']

        # Can be calcuated for rouge1, rouge2, rougeL, rougeLsum
        # precision = precision.item()
        # recall = recall.item()
        fmeasure = fmeasure.item()

        # return {
        #     self.name : {'precision' : {'value' : precision},
        #                   'recall' : {'value': recall},
        #                     'fmeasure' :{'value':  fmeasure}
        #     }
        # }
        return {
            self.name : round(fmeasure,2)
        }

    
    def risk(self,score):
        threshold = 0.7
        # precision_threshold = 'pass' if score[self.name]['precision'] >= threshold else 'fail'
        # recall_threshold = 'pass'  if score[self.name]['recall'] >= threshold else 'fail'
        # fmeasure_threshold  = 'pass' if score[self.name]['fmeasure'] >= threshold else 'fail'

        return {
            self.name: "pass" if score[self.name] >= threshold else "fail"
        }