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
from qa_squad import *
import requests

def get_vllm_llama(temperature, max_tokens, chats):
    url = "http://127.0.0.1:60000/ask"

    data = {
        "temperature": temperature,
        "max_tokens": max_tokens,
        "chats": chats,
    }
    
    response = requests.post(url, json=data, timeout=None)
    response_data = response.json()

    passage_ls = response_data.get("output", [])
    return passage_ls

def call_llama3_single_prompt(
    inputs, model="Llama-3.1-8B-Instruct", max_decode_steps=20, temperature=0.8
):
      
    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)
        
    return ans


def call_llama3_func(
    inputs, model="Llama-3.1-8B-Instruct", max_decode_steps=1024, 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
    )
    
    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])},
            {"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 cover_em(pred_ls, ans_ls):
    assert len(pred_ls) == len(ans_ls)
    cnt = 0
    for idx in range(len(pred_ls)):
        pred = pred_ls[idx].lower()
        ans = eval(ans_ls[idx])
        for j in range(len(ans)):
            if ans[j].lower() in pred:
                cnt += 1
                break
                
    return cnt/len(pred_ls)



def supervised_method():
    sub_query_str_l = get_supervised_decom(raw_queries)
    passages_ls = get_ir_result(raw_queries, sub_query_str_l)
    
    pred_ls = get_ans(raw_queries, passages_ls)
    
    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']])
        
    passages_ls = get_ir_result(raw_queries, sub_query_str_l)
    
    pred_ls = get_ans(raw_queries, passages_ls)
    
    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']])
        
    passages_ls = get_ir_result(raw_queries, sub_query_str_l)
    
    pred_ls = get_ans(raw_queries, passages_ls)
    
    return cover_em(pred_ls, true_answers)



if __name__ == '__main__':
    dataset_path = ragqa_paths.dataset_file("squad", "squad_test.csv")
    tokenizer = AutoTokenizer.from_pretrained(ragqa_paths.LLAMA_MODEL)
    client = OpenAI(api_key="0",base_url="http://0.0.0.0:50001/v1")

    raw_data = pd.read_csv(dataset_path, header=0)
    
    # raw_data = raw_data.head(5)
    
    raw_queries = list(raw_data['question'])
    true_answers = list(raw_data['answers'])
    
    print(supervised_method())
    print(unsupervised_method())
    print(iclfeed_method())