from challenge_gener import Dynamic_Intelligent_Assessment
from DIA_benchmark_evalutor import LLMChallengeEvaluator
import pandas as pd
import json
import os
import re
from tqdm import tqdm 

TEMPLATE_NUM = 122

# hyper para
K = 5
MODEL_NAME = 'hunyuan-turbo'
URL = ''
KEY = ''


# get generated challenge data
generator = Dynamic_Intelligent_Assessment()
challenge_data = generator.generate_json(k=K)

# create model client 
target_model = LLMChallengeEvaluator(host=URL, api_key=KEY, model=MODEL_NAME)

# assert the data type is json type
if isinstance(challenge_data, str):
    json_data = json.loads(challenge_data)

with tqdm(total=K * TEMPLATE_NUM, desc="Eval " + str(MODEL_NAME), unit="question") as pbar:
    for i, item in enumerate(json_data['questions']):
        # infer
        q_id = str(item['challenge']['template_id']) + '-' + str(item['challenge']['instance'])
        template_id = str(item['challenge']['template_id'])

        description = item['challenge']['description']
        instructions = item['challenge']['instructions']
        q = f"{description}\n{instructions}"    
        llm_response = target_model.ask_llm(q)
        res = target_model.clean_and_extract_xml(llm_response) if llm_response else None

        # eval
        correct_solution = item['solution']['challenge_solution']

        if llm_response == correct_solution:
            status = "Correct"
        elif llm_response == "<xml>I-DO-NOT-KNOW</xml>":
            status = "Skipped"
        else:
            status = "Incorrect" 



        ILLEGAL_CHARACTERS_RE = re.compile(r'[\000-\010]|[\013-\014]|[\016-\037]')
        q = ILLEGAL_CHARACTERS_RE.sub(r'', q)
        llm_response = ILLEGAL_CHARACTERS_RE.sub(r'', llm_response)

        save_data = {
            'question_id': q_id,
            'template_id': template_id,
            'question': q,
            MODEL_NAME + '_response': llm_response,
            'status': status,
            'reference': correct_solution
        }

        # 把json数据保存为excel格式
        file_path = MODEL_NAME + '_DIA-' + str(K) +'_results.xlsx'
        if not os.path.exists(file_path):
            df = pd.DataFrame(save_data, index=[0])
            df.to_excel(file_path, index=False)
        else:
            df = pd.read_excel(file_path)
            new_df = pd.DataFrame([save_data])
            df = pd.concat([df, new_df], ignore_index=True)
            df.to_excel(file_path, index=False)

        pbar.update(1)