import time
import os
import json

import utils.utils as utils
import utils.gpt as gpt
import retrieval.solve_retrieval as solve_retrieval


def solve(data, index, client, logger, args):
    """ solve the problem calling GPT """
    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, solution, answer = thy.get("informal problem", ""), thy.get("solution", ""), thy.get("answer", "")
    if (solution != "" or informal_problem == "") and args.overwrite == 0:
        logger.info('skip')
    else:
        informal_problem = thy.get("informal problem", "")
        t0 = time.time()
        # prompt = f"Please think step by step and give me the answer after \"The answer is\". {informal_problem}"
        prompt = f"Solve the following math problem. Give me detailed solution steps and write the final numer in the sentence \"The answer is \". {informal_problem}"
        solve_examples = solve_retrieval.prob_sample(prompt, k=args.num_retrive)
        natural_solution = gpt.gpt4_response(client, prompt, solve_examples, top_p=0.95)
        logger.info(informal_problem + '\nTO ========>>>\n' + natural_solution + "\nThe correct answer is: " + answer)
        t1 = time.time()
        logger.info('success! In (%s)s' % (t1-t0))
        thy['solution'] = natural_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 