import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
os.environ['HTTP_PROXY'] = 'http://127.0.0.1:7890'
os.environ['HTTPS_PROXY'] = 'http://127.0.0.1:7890'

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'))
import ragqa_paths  # [ragqa] portable paths
from vllm import LLM, SamplingParams
from transformers import AutoTokenizer
from openai import OpenAI

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

from multiqa import *
import requests
import pickle
from tqdm import tqdm
# from byaldi import RAGMultiModalModel

dataset = 'multiqa'


# colpali = RAGMultiModalModel.from_index(dataset, index_root = "/data1/liuyaoyang/Papers/icml2025/multi_rag/byaldi/indexes")


import re

def filter_subqueries(data, queries):
    filtered_data = []
    
    for group, query in zip(data, queries):
        filtered_group = []
        query_tokens = set(re.findall(r"\w+", query.lower()))  # 提取queries中的单词
        
        for subquery in group[0]:
            subquery_tokens = re.findall(r"\w+", subquery)  # 提取subquery中的单词
            filtered_subquery = " ".join([token for token in subquery_tokens if token.lower() in query_tokens])
            filtered_group.append(filtered_subquery)
        
        filtered_data.append([filtered_group])
    
    return filtered_data

def call_llama3_single_prompt(
    inputs, model="Llama-3.1-8B-Instruct", max_decode_steps=20, temperature=0.0
):
      
    inputs_ls = []
    if isinstance(inputs, str):
        messages = [
            {"role": "user", "content": inputs},
        ]
        inputs_ls.append(tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True))

    else:
      for idx in range(len(inputs)):
          inputs_ls.append(tokenizer.apply_chat_template(inputs[idx], tokenize=False, add_generation_prompt=True))
    
    # ans = get_vllm_llama(temperature, max_decode_steps, inputs_ls)
    results = client.completions.create(
            model=ragqa_paths.LLAMA_MODEL,
            max_tokens=max_decode_steps,
            temperature=0,
            prompt=inputs_ls,
            timeout = None
    )
    ans = []
    for item in results.choices:
        ans.append([item.text.strip()])
        
    return ans


def call_llama3_func(
    inputs, model="Llama-3.1-8B-Instruct", max_decode_steps=100, temperature=0.0
):

  print(max_decode_steps, temperature)
  
  output = call_llama3_single_prompt(
      inputs,
      model=model,
      max_decode_steps=max_decode_steps,
      temperature=temperature
  )

  if isinstance(inputs, str):
    return output[0]
  else:
    return output
      

def get_supervised_decom(queries_ls):
    prompts = []
    for query in queries_ls:
        prompts.append([{"role": "user", "content": query}])
        
    prompts_tokened = [tokenizer.apply_chat_template(x, tokenize=False, add_generation_prompt=True) for x in prompts]
    
    results = client.completions.create(
            model="supervised",
            max_tokens=512,
            temperature=0,
            prompt=prompts_tokened,
            timeout = None
    )
    
    ans = []
    for item in results.choices:
        ans.append([[x.strip() for x in item.text.split('|')]])
    
    return ans

def get_ans(queries_ls, passages):
    prompts = []
    for i in range(len(queries_ls)):
        prompts.append([
            {"role": "system", "content": "You are a helpful assistant. Please answer the question to the best of your knowledge, even if the context does not directly provide the information. Use any relevant knowledge you have to provide a helpful answer."},
            {"role": "user", "content": 'Context: ' + '\n'.join(passages[i][:3])},
            {"role": "user", "content": 'Question: ' + queries_ls[i]},
        ])
        
        
    pred_ls = [row[0] for row in call_llama3_func(prompts, max_decode_steps=100)]
    
    return pred_ls



def hit_score(passages_ls, anspids, k):
    assert len(passages_ls) == len(anspids)

    cnt = 0
    for i in range(len(passages_ls)):
        retrieved_100_topk = passages_ls[i][:k]
        ans_raw = anspids[i]
        for ans in ans_raw:
            if ans in retrieved_100_topk:
                cnt += 1
                break
    return cnt / len(passages_ls)
    

def cover_em(predictions, ground_truths):
    score = []
    for i in range(len(predictions)):
        pred = predictions[i].lower().strip()
        gt_ls = ground_truths[i].lower().split(',')
        for gt in gt_ls:
            if gt in pred:
                score.append(1)
                break
        else:
            score.append(0)
    return sum(score) / len(score)




def supervised_method():
    sub_query_str_l = get_supervised_decom(raw_queries)

    sub_query_str_l = filter_subqueries(sub_query_str_l, raw_queries)

    pred_ls, re_imgs= get_eval_answer_llava(raw_queries, sub_query_str_l, ans_pids, patch_emb_by_img_ls)

    detailed_results_df = pd.DataFrame(
      list(
          zip(
              raw_queries,
              sub_query_str_l,
              true_answers,
              pred_ls,
              re_imgs,
              
          )
      ),
      columns=[
          "raw_queries",
          "sub_queries_ls",
          "true_answer",
          "pred_answers",
          "re_imgs",
      ],
    )
    detailed_results_df.to_csv('supervised_method.csv')

    return cover_em(pred_ls, true_answers)


def unsupervised_method():
    sub_query_str_l = []
    for query in tqdm(raw_queries):
        url = 'http://127.0.0.1:50002/execute?query='+query
        response = requests.get(url=url)
        res_dic = response.json()
        sub_query_str_l.append([res_dic['text']])
    
    sub_query_str_l = filter_subqueries(sub_query_str_l, raw_queries)
        
    # print(sub_query_str_l)
    try:
        pred_ls, re_imgs= get_eval_answer_llava(raw_queries, sub_query_str_l, ans_pids, patch_emb_by_img_ls)
    except:
        pass

    detailed_results_df = pd.DataFrame(
      list(
          zip(
              raw_queries,
              sub_query_str_l,
              true_answers,
              pred_ls,
              re_imgs,
              
          )
      ),
      columns=[
          "raw_queries",
          "sub_queries_ls",
          "true_answer",
          "pred_answers",
          "re_imgs",
      ],
    )
    detailed_results_df.to_csv('unsupervised_method.csv')
    
    return cover_em(pred_ls, true_answers)
    


def iclfeed_method():
    sub_query_str_l = []
    
    for query in tqdm(raw_queries):
        url = 'http://127.0.0.1:50003/execute?query='+query
        response = requests.get(url=url)
        res_dic = response.json()
        sub_query_str_l.append([res_dic['text']])
        
    sub_query_str_l = filter_subqueries(sub_query_str_l, raw_queries)

    try:
        pred_ls, re_imgs= get_eval_answer_llava(raw_queries, sub_query_str_l, ans_pids, patch_emb_by_img_ls)
    except:
        pass

    detailed_results_df = pd.DataFrame(
      list(
          zip(
              raw_queries,
              sub_query_str_l,
              true_answers,
              pred_ls,
              re_imgs,
              
          )
      ),
      columns=[
          "raw_queries",
          "sub_queries_ls",
          "true_answer",
          "pred_answers",
          "re_imgs",
      ],
    )
    detailed_results_df.to_csv('iclfeed_method.csv')
    
    return cover_em(pred_ls, true_answers)

   
def dense_method():
    sub_query_str_l = [[[raw]]for raw in raw_queries]
    try:
        pred_ls, re_imgs= get_eval_answer_llava(raw_queries, sub_query_str_l, ans_pids, patch_emb_by_img_ls)
    except:
        pass

    detailed_results_df = pd.DataFrame(
      list(
          zip(
              raw_queries,
              sub_query_str_l,
              true_answers,
              pred_ls,
              re_imgs,
              
          )
      ),
      columns=[
          "raw_queries",
          "sub_queries_ls",
          "true_answer",
          "pred_answers",
          "re_imgs",
      ],
    )
    detailed_results_df.to_csv('dense_method.csv')

    return cover_em(pred_ls, true_answers)

def colbert_method():
    tmp = gen_prompt()
    
    raw_op_prompts = call_llama3_func(tmp)
    
    # print(raw_op_prompts[:3])
    
    sub_queries_ls= []
    for idx in range(len(raw_queries)):
        tmp_ls = raw_op_prompts[idx][0].replace("\n", "").split(",")
        tmp_ls = [list(set([item.strip() for item in tmp_ls if item.strip() and item.strip() in raw_queries[idx]]))]
        if len(tmp_ls[0]) == 0:
            tmp_ls = [[raw_queries[idx]]]
        sub_queries_ls.append(tmp_ls)

    try:
        pred_ls, re_imgs= get_eval_answer_llava(raw_queries, sub_queries_ls, ans_pids, patch_emb_by_img_ls)
    except:
        pass
    
    
    detailed_results_df = pd.DataFrame(
      list(
          zip(
              raw_queries,
              sub_queries_ls,
              true_answers,
              pred_ls,
              re_imgs,
              
          )
      ),
      columns=[
          "raw_queries",
          "sub_queries_ls",
          "true_answer",
          "pred_answers",
          "re_imgs",
      ],
    )
    detailed_results_df.to_csv('colbert_result.csv')


    return cover_em(pred_ls, true_answers)

def colpali_method():
    re_img_ls = []
    
    for query in tqdm(raw_queries):
        results = colpali.search(query, k=100)
        re_img_ls.append([x['metadata'][0]['filename'] for x in results])
        
    pred_ls = colbert_score(raw_queries, re_img_ls, ans_pids, dataset)
    return cover_em(pred_ls, true_answers)
    


def gen_prompt():
    
    prompts = []
    for query in raw_queries:
        prompt = []
        prompt.append({"role": "system", "content": """"Given the input query, break it down into meaningful tokens like ColBERT. Ensure the tokens retain the key semantic components. Provide the output as a comma-separated list. Query: '{query}' Tokens:"""})
        prompt.append({"role": "user", "content": "Query: Victoria Hong Kong has many what type of buildings?"})
        prompt.append({"role": "assistant", "content": "Victoria, Hong, Kong, has, many, what, type, of, buildings,?"})

        prompt.append({"role": "user", "content": f"Query: {query}" })
        prompts.append(prompt)

    return prompts


def without_method():
        
    pred_ls = wo_llava_vllm(raw_queries)
    
    return cover_em(pred_ls, true_answers)


if __name__ == '__main__':
    dataset_path = ragqa_paths.dataset_file(dataset, f"{dataset}_test.csv")
    tokenizer = AutoTokenizer.from_pretrained(ragqa_paths.LLAMA_MODEL)
    client = OpenAI(api_key="0",base_url="http://127.0.0.1:50001/v1")

    raw_data = pd.read_csv(dataset_path, header=0)
    raw_data = raw_data.drop_duplicates(subset=['question'])

    
    raw_queries = list(raw_data['question'])
    true_answers = list(raw_data['answer'])
    ans_pids = list(raw_data['image'])

  
    
    # print(supervised_method())
    # print(unsupervised_method())
    # print(iclfeed_method())
    # print(colbert_method())
    # print(colpali_method())
    print(dense_method())
    # print(without_method())