"""
The Rouge Evaluator is an advanced AI tool designed to assess the quality of dialogue models. 
It uses a unique approach to evaluate the responses generated by these models, focusing on aspects 
such as relevance, coherence, and fluency.

This tool is particularly useful for developers and researchers working on dialogue systems, as it
allows them to measure the effectiveness of their models and make necessary enhancements. 

The Rouge Evaluator is a valuable asset for anyone looking to enhance the quality 
and performance of their dialogue models.
"""

import rouge  # type:ignore

from ..schemas.common_structures import InputData
from ..schemas.evaluator_config import EvaluatorOutput, RougeEvaluatorConfig
from ..schemas.experiment_config import ExperimentResult, MultimodalOutput
from .base_evaluator import BaseEvaluator


class RougeEvaluator(BaseEvaluator):
    """Evaluator using rouge to calculate rouge score"""

    default_config = RougeEvaluatorConfig(name="rouge_evaluator")

    def __init__(self, config: RougeEvaluatorConfig):
        super().__init__(config)
        self.config = config
        self.rough = rouge.Rouge()

    def evaluate(self, experiment_result: ExperimentResult) -> EvaluatorOutput:
        """Evaluate the experiment result using rouge evaluat"""
        assert isinstance(self.config, RougeEvaluatorConfig)
        input_data = experiment_result.input_data
        raw_output = experiment_result.raw_output.text_output
        expected_result = input_data.expected_result

        scores = self.rough.get_scores(
            hyps=raw_output, refs=expected_result, avg=True
        )

        return EvaluatorOutput(
            name=self.config.name,
            display_name="rouge",
            result=scores,
            metric_calculators=self.config.metric_calculators
        )


BaseEvaluator.register_evaluator(
    "rouge_evaluator", RougeEvaluator, RougeEvaluatorConfig
)


def main():
    """Main function to test the rouge evaluator"""
    evaluator_config = RougeEvaluatorConfig(
        name="rouge_evaluator",
        display_name="rouge",
        metric_calculators=[],
    )
    input_data_example = InputData(
        content={
            "instruction": "translate the sentence to english",
        },
        expected_result="Have a great day!"
    )
    experiment_result_example = ExperimentResult(
        input_data=input_data_example,
        combination={
            "wrapper1": "var1",
            "wrapper2": "var2"
        },
        raw_output=MultimodalOutput(text_output="have a nice day"),
        latency=30.0,
        token_usage=20
    )

    evaluator = RougeEvaluator(evaluator_config)
    result = evaluator.evaluate(experiment_result_example)
    print(result)


if __name__ == "__main__":
    main()