# -*- coding: utf-8 -*-
# This script is used to evaluate the performance of Pigeon AI Video Translation system by using Large Language Model.

# Written by Jiaen LIU, 2023/09/18

# Import the necessary packages
from langchain.evaluation import load_evaluator, EvaluatorType
from langchain.prompts import PromptTemplate
from langchain.chat_models import ChatOpenAI
# from src.srt_util.srt import SrtScript

# Load the evaluator

def init_evaluator():

    llm = ChatOpenAI(temperature=0, model="gpt-4-0613")

    fstring = """You are an expert English to Chinese translator specialized in Startcraft2.
            You are grading the following question:
            {input}
            Here is the real answer:
            {reference}
            You are grading the following predicted answer:
            {output}
            based on the following criteria: 
            {criteria}
            Give two grades, one for completness and another for accuracy and rate them from a scale of 0 to 100, where 0 is the lowest (very low completeness/accuracy) and 100 is the highest (very high completness/accuracy)? 
            Do not base the two scores off each other give them the scores independently. Give explanations for every single one and if the answer if partially correct that is acceptable. However punish the scores for answers that are 
            numerically incorrect this also includes values that have the $ in front
            Please give the completeness score first followed by the accuracy score. 
            For example: Completeness: 70. Accuracy: 40. Explanation here
            Do not differ from the format ever
            """
    prompt = PromptTemplate.from_template(fstring)

    return load_evaluator("labeled_criteria", llm=llm, prompt=prompt, criteria="correctness")

def evaluate_prediction(input, reference, prediction, evaluator):
    eval_result = evaluator.evaluate_strings(
        prediction=prediction,
        input=input,
        reference=reference,
    )
    return eval_result

if __name__ == "__main__":
    evaluator = init_evaluator()
    eval_result = evaluate_prediction("this is an test sentences", "这不是一个测试语句。", "这是一个测试句子。", evaluator)
    print(eval_result)