import time
import os
import json

import utils.utils as utils
import utils.gpt as gpt
import retrieval.inauto_retrieval as inauto_retrieval

def informalize(data, index, client, logger, args):
    """ informalize the problem """
    Q, A = data['question'], data['answer']
    file_name = "problem_%s" %(args.index[0] + index)
    save_file_path = os.path.join(args.save_folder_path, file_name+'.json')
    if not os.path.exists(save_file_path):
        answer = utils.parse_answer(A)
        thy_series = {'problem': Q, 
                    'solution': A, 
                    'answer': answer}
        with open(save_file_path, 'w') as file:  
            json.dump(thy_series, file, indent=4)
    else:
        with open(save_file_path) as file:
            thy_series = json.load(file)
            answer = thy_series['answer']     
    thy = thy_series.get(f'generation_{max(args.save_iter,0)}', {})
    informal_problem, answer = thy.get("informal problem", ""), thy.get("answer", "")
    if (informal_problem != "" and answer != "" and "SMT-LIB" not in informal_problem) \
                                and args.overwrite == 0:
        logger.info('skip')
    else:
        formal_problem = thy.get("formal problem", "")
        t0 = time.time()
        if args.informal_prompt == "origin":
            prompt = f"Translate the math problem formulated with SMT-LIB back to a natural language problem. {formal_problem}"
        elif args.informal_prompt == "math-word":
            prompt = f"Translate the math problem formulated with SMT-LIB back to a natural language problem. Please ensure the problem to be a math word problem. {formal_problem}"
        elif args.informal_prompt == "modify":
            last_thy = thy_series[f"generation_{args.ref_iter}"]
            origin_problem = last_thy["formal problem"]
            prompt = f"Modify the given math word problem {origin_problem}, such that it safisfies the following SMT-LIB formulas {formal_problem}"
        prob_examples = inauto_retrieval.prob_sample(prompt, k=args.num_retrive)
        natural_problem = gpt.gpt4_response(client, prompt, prob_examples, top_p=0.95)
        logger.info(formal_problem + '\nTO ========>>>\n' + natural_problem)
        t1 = time.time()
        logger.info('success! In (%s)s' % (t1-t0))
        thy["informal problem"] = natural_problem
        thy["solution"] = "" # clean the solution
        thy_series.update({f'generation_{max(args.save_iter,0)}' : thy})
        with open(save_file_path, 'w') as file:  
            json.dump(thy_series, file, indent=4)
    return