from transformers import pipeline
import torch
import pandas as pd
import sys
import os
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..', '..', 'document_retrieval', 'Decompose_retrieval'))
from pipline import *



def ask_gpt(query):
    messages = [{"role": "user", "content": query}]
    result = client.chat.completions.create(messages=messages, model="meta-llama/Meta-Llama-3-8B-Instruct")
    return result.choices[0].message.content


def get_accuracy_re(true_imgs, re_imgs):
    # print(true_imgs)
    # print(re_imgs)
    cnt = 0
    for i in range(len(true_imgs)):
        i_gt = true_imgs[i].rsplit('.', 1)[0]
        image_retrieval = re_imgs[i].rsplit('.', 1)[0]
        if i_gt == image_retrieval:
            cnt += 1
    return cnt/len(true_imgs)


def get_rag_answers(q_list, sub_q_list, dataset_name = "multiqa"):
    if dataset_name == "multiqa":
        ans_img_list, re_im = eval_acc(q_list, sub_q_list, patch_emb_by_img_ls)
        ans = []
        re_imgs = []
        for i in range(len(ans_img_list)):
            ans.append(ans_img_list[i][0])
            re_imgs.append(re_im[i])
        return ans, re_imgs
    

def exact_match(predictions, ground_truths):
    # sum = 0
    score = []
    for i in range(len(predictions)):
        pred = predictions[i].lower()
        gt = ground_truths[i].lower()
        if gt == pred:
            score.append(1)
        else:
            score.append(0)
    return score


def get_accuracy_multiqa(predictions, ground_truths):
    score = exact_match(predictions, ground_truths)
    return score

dataset_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "multiqa_test.csv")
client = OpenAI(api_key="0",base_url="http://0.0.0.0:8000/v1")

single_task_df = pd.read_csv(dataset_path, header=0)
raw_data = single_task_df[['query', 'answer', 'image']]
raw_queries = list(raw_data['query'])
true_answers = list(raw_data['answer'])
true_images = list(raw_data['image'])


raw_op_prompts = []
for i in range(len(raw_data)):
    raw_pred = ask_gpt(raw_queries[i])
    raw_op_prompts.append(raw_pred)


pred_answers, re_imgs = get_rag_answers(raw_queries, raw_op_prompts, "multiqa")


accuracy = get_accuracy_multiqa(
  predictions = pred_answers,
  ground_truths = true_answers
)
re_acc = get_accuracy_re(true_images , re_imgs)

accuracy = np.average(accuracy)

print(accuracy, re_acc)