import json
import os
import random
from tqdm import tqdm
import openai
import datetime
import time
import argparse
import re

openai.api_key=''
openai.organization = ""
openai.Model.list()
SYS_VQA="""
Please act as an impartial judge and conduct a comprehensive assessment of a multimodal AI assistant's performance in the field of Visual Question Answering (VQA). Each data sample to be evaluated follows the following format:
[Question]
{Question}\n
[Ground Truth]
{Ground truth}\n
[Assistant's Chain of Thought]
{Chain of thought generated by AI assistant}\n
[Assistant's Final Choice]
{Final Choice generated by AI assistant}\n
Your task is to evaluate the quality of natural language generation from AI assistant considering factors such as the helpfulness, relevance, accuracy, depth, creativity, and level of detail of the response. 
Please first provide a comprehensive explanation of your evaluation and an overall score ranging from 0 to 10 based on explanation, where a higher score indicates better overall performance. Please output in the following format:\n
[Explanation] {Evaluation Explanation}
[Overall Score] {An integer ranging from 0 to 10 representing the final evaluation score}\n
Please ensure that your evaluation score comprehensively captures the AI assistant's performance avoiding any potential bias. Assuming that the visual information mentioned by the AI assistant is contained in the image, you only need to evaluate the quality of the generated text. Your assessments will contribute to enhancing the assistant's effectiveness in visual question answering.
"""


DATA_Template = """
[Question]
{}\n
[Ground Truth]
{}\n
[Assistant's Chain of Thought]
{}\n
[Assistant's Final Choice]
{}\n
"""

def clean_question(question): # delete context
    qlist = question.split('Options:')
    q = qlist[0].split('Context:')
    return 'Question: ' + q[0] + 'Options:' + qlist[1] + "\n"

def option_text(id):
    option = ['A','B','C','D','E','F','G']
    return '(' + option[id] + ') ' 

def parse_model_name(answer_path):
    js_name = os.path.basename(answer_path)
    match = re.match(r'([^_]+)', js_name)
    return match.group(1)

def generate_query(gt_data, answer):
    question = clean_question(gt_data['query'])
    gt_choice = gt_data['gt_choice']
    gt = option_text(gt_choice)
    Cot = answer['CoT_answer']
    idx = Cot.rfind('.')
    if idx>0:
        Cot = Cot[:idx+1]
    answer_text = answer['answer']+')'
    query = DATA_Template.format(question, gt, Cot, answer_text)
    return query


def parse_score(text):
    #match = re.search(r'\[Overall Score\]: (\d+)', text)
    match = re.search(r'\[Overall Score\]\s*(\d+(?:\.\d+)?)',text)
    if match:
        score = float(match.group(1))
        #print("Overall Score:", score)
    else:
        print(f"Overall Score not found. origin text: {text}")
        score=None
    return score

def GPT_Metric(base_data_path, answer_path, response_dir):
    dataset = json.load(open(base_data_path,'rb'))
    answer = json.load(open(answer_path,'rb'))
    model_name = parse_model_name(answer_path)[:-5]
    cur_time = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    save_base_dir = os.path.join(response_dir, model_name, cur_time)
    os.makedirs(save_base_dir, exist_ok=True)
    response_path = os.path.join(save_base_dir, 'sqa_gpt_metric.jsonl')
    res_path = os.path.join(save_base_dir, 'results.jsonl')
    ans_file = open(response_path,'w')
    res_file = open(res_path, 'w')
    rd_seed=0
    random.seed(rd_seed)
    sample_ids = random.sample(range(2017), 250)
    sample_ids.sort()
    print(f'Preparing queries ...')
    query_list=[]
    for id in tqdm(sample_ids):
        query = generate_query(dataset[id],answer[id])
        query_list.append(query)
    print('GPT evaluating ...')
    response_list = []
    score_list = []
    res_list=[]
    
    for ids,query in tqdm(zip(sample_ids, query_list)):
        messages = [{"role":"system", "content": SYS_VQA}]
        messages.append({"role":"user", "content":query})
        while True:
            try:
                response = openai.ChatCompletion.create(model="gpt-4", messages=messages,n=6, max_tokens=512)
                response_list.append(response)
                break
            except:
                continue
        cur_text = []
        cur_score = []
        for ans in response['choices']:
            cur_text.append(ans['message']['content'])
            score = parse_score(ans['message']['content'])
            if score!=None:
                cur_score.append(score)
        input_tokens=response["usage"]["prompt_tokens"]
        comp_tokens=response["usage"]["completion_tokens"]
        cost = input_tokens*0.03/1000+comp_tokens*0.06/1000
        cur_dict={
            'ids':ids,
            'query':query,
            'response': cur_text,
            'scores': cur_score,
            'cost': cost,
            'origin_answer':answer[ids],

        }
        res_list.append(cur_dict)
        ans_file.write(json.dumps(response) + "\n")
        res_file.write(json.dumps(cur_dict)+"\n")
        ans_file.flush()
        res_file.flush()
        time.sleep(2)

    with open(os.path.join(save_base_dir, 'samples_results.json'), 'w', encoding='utf-8') as f:
            f.write(json.dumps(res_list, indent=4))
    total_score, total_cost = 0,0
    for res in res_list:
        tmp=0
        for i in res['scores']:
            tmp+=i
        tmp/=len(res['scores'])
        total_score+=tmp
        total_cost+=res['cost']
    total_score/=len(res_list)
    print(f"GPT Evaluation on {model_name} completed. Avg:{total_score}, Cost:{total_cost}")
    overall_res={
        'model name':model_name,
        'Avg Score': total_score,
        'Cost':total_cost
    }
    with open(os.path.join(save_base_dir, 'results.json'), 'w', encoding='utf-8') as f:
            f.write(json.dumps(overall_res, indent=4))
    return response_list, res_list
    
def eval_score(response_list):
    total = 0
    cnt = 0
    for item in response_list:
        text = item['choices'][0]['message']['content']
        score = parse_score(text)
        if score!=None and score>=0:
            total+=score
            cnt+=1
    return total/cnt*10.0


def eval_cost(response_list):
    inputs=0
    outputs = 0
    cost=0.0
    for item in response_list:
        ic= item["usage"]["prompt_tokens"]
        oc = item["usage"]["completion_tokens"]
        inputs+=ic
        outputs+=oc
    print(inputs)
    print(outputs)
    cost = inputs/1000*0.0015+outputs/1000*0.002
    return cost

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--base-data-path', default='../data/LAMM/2D_Benchmark/meta_file/VQA_ScienceQA.json')
    parser.add_argument('--answer-path', type=str, required=True)
    parser.add_argument('--response-dir', default='../results/gpt_eval')
    args = parser.parse_args()
    response_list, res_list = GPT_Metric(args.base_data_path,
                               args.answer_path,
                             args.response_dir)