import json, pandas as pd, tqdm
from Scripts.Prompting import *
from Scripts.Utils import *

ip = r"Responses\cache\old_info\1_mbpp-updated_corrected_50_ambiguous_0.json"
# update for new information
with open(ip, 'r') as f:
    data = json.load(f)
args=get_args(None)
args.info_prompt = "agg_comp_imp_fewshot_updated"
for d in tqdm.tqdm(data):
    d['info'] = get_info(d['text'], args, None, 0)[0]
    with open(ip.replace('old_info', 'new_info_agg'), 'w') as f:
        json.dump(data, f, indent=4)


# try:
#     with open(ip.replace('_corrected','_corrected_rephrased_3'), 'r') as f:
#         data = json.load(f)
# except:
#     with open(ip, 'r') as f:
#         data = json.load(f)

# prompt_file = r"Scripts\config_experiment\prompts\rephrasingNL.md"
# prompt_file = r"Scripts\config_experiment\prompts\info2NL.md"
# prompt_file = r"Scripts\config_experiment\prompts\rephrase_info2NL.md"
# with open(prompt_file, 'r') as f:
#     prompts = f.read()

# def rephrase(text, info):
#     # prompt = prompts.replace('{{originalNL}}', text)
#     prompt = prompts.replace('{{info}}', str(info))
#     text = get_llm_response('gpt4-turbo', prompt, stop=['---'])[0]
#     # for s in ['`', '---', 'query', 'info']:
#     #     text = text.split(s)[0].strip()
#     text = text.split('`')[2].split(":")[1].split('\n')[0].strip()
#     return text

# for d in tqdm.tqdm(data):
#     if "rephrased" in d: continue
#     d['text'] = rephrase(d['text'], d['info'])
#     d['rephrased'] = True
#     with open(ip.replace('_corrected','_corrected_rephrased_3'), 'w') as f:
#         json.dump(data, f, indent=4)


# df = pd.DataFrame(data)
# df.to_csv(ip.replace('_corrected','_corrected_rephrased_3').replace('.json', '.csv'), index=False)
