import argparse
from tot.methods.bfs import solve, solve_new
from tot.tasks.game24 import Game24Task
from tot.tasks.blocksworld2 import Blocksworld
import sys
from utils import *
from prompts import *
import copy
from tqdm import tqdm

model_name="gpt-3.5-turbo-0301"
args = argparse.Namespace(backend=model_name, temperature=1.0, \
                          task='blocksworld', naive_run=False, prompt_sample=None, \
                          method_generate='propose', method_evaluate='value', \
                          method_select='greedy', n_generate_sample=1, n_evaluate_sample=3, \
                          n_select_sample=5)

mm = {2: 30, 4: 56, 6: 114}

for _ in range(5):    
    for j in [2, 4, 6]:
        res = []
        for idx in range(mm[j]):
            task = Blocksworld(step_count=j)
            ys, infos = solve_new(args, task, idx)

            for ys_ in ys:
                try:
                    gt_action_sequence = [i.split(':')[1].strip() for i in ys_.split('\n') if len(i)>5]
                except:
                    gt_action_sequence = []
                item = task.item
                
                init_block_config = state_text2json(item['real_problem'], item['participating_blocks'])
                final_block_config = copy.deepcopy(init_block_config)
                for action in gt_action_sequence:
                    final_block_config, valid_action = add_action_to_json_state(final_block_config, action)

                if final_block_config == task.final_block_config:
                    res.append(1)
                    break
        print('Result:\t', sum(res), j)
