from LLM import *
from prompt_env3 import *
from env3_create import *
from sre_constants import error
import random
import os
import json
import re
import copy
import numpy as np
import shutil
import time
from env3_func import *
import tiktoken
from argparse import ArgumentParser

parser = ArgumentParser()
parser.add_argument('-experiment_trial_num', '--experiment_trial_num', type=int, default=1) # experiemnt trial number
parser.add_argument('-input_error_prompt_token_limit', '--input_error_prompt_token_limit', type=int, default=15000) # The token limit for the input error prompt
parser.add_argument('-model_name_promptLLM', '--model_name_promptLLM', default='gpt-4-1106-preview') # The model name for the promptLLM
parser.add_argument('-model_name_testLLM', '--model_name_testLLM', default='gpt-3.5-turbo-16k-0613') # The model name for the testLLM
parser.add_argument('-min_level', '--min_level', type=int, default=2) # The minimum levels of evolution for the prompt optimization

parser.add_argument('-n_children', '--n_children', type=int, default=8) # The number of son prompts to be expanded in each level
parser.add_argument('-n_selected', '--n_selected', type=int, default=2) # The number of best prompts for further optimization in each level

parser.add_argument('-prompt_method', '--prompt_method', default='PROMST') # The prompt optimization method
parser.add_argument('-with_score_model', '--with_score_model', default='False') # Whether to use the score model
parser.add_argument('-Training_path', '--Training_path', default='../BoxLift/train_set/') # The path to the training set
parser.add_argument('-Testing_path', '--Testing_path', default='../BoxLift/test_set/') # The path to the testing set
parser.add_argument('-base_path', '--base_path', default='../BoxLift/') # The base path for the training set

args = parser.parse_args()
experiment_trial_num = args.experiment_trial_num
input_error_prompt_token_limit = args.input_error_prompt_token_limit
model_name_promptLLM = args.model_name_promptLLM
model_name_testLLM = args.model_name_testLLM
min_level = args.min_level
prompt_method = args.prompt_method
with_score_model = args.with_score_model
Testing_path = args.Testing_path
Training_path = args.Training_path
base_path = args.base_path
n_children = args.n_children
n_selected = args.n_selected

enc = tiktoken.get_encoding("cl100k_base")
assert enc.decode(enc.encode("hello world")) == "hello world"
enc = tiktoken.encoding_for_model("gpt-4")

# Create the training and testing environments
if not os.path.exists(Testing_path) or not os.path.exists(Training_path):
    create_env3(Testing_path, repeat_num = 4)
    create_env3(Training_path, repeat_num = 4)

Starting_prompt_task_explain = f'''
  You are a central planner directing lifting agents in a warehouse to lift boxes. Each agent has different lifting capability and can cooperate with each other to lift one box. In summation of lifting capability, the agents can lift all boxes. 

  The boxes are identified by their volume, e.g., box[1.4V]. The agents are identified by their lifting weight capability, e.g., agent[1.5W]. Actions are like: "box[1.7V]":"agent[2.5W]", "box[6.0V]":"agent[1.5W], agent[2.5W]".

  Your task is to divide the group of each agent to lift all the boxes. After each step, environments provide updates for the left boxes. Your job is to coordinate the agents optimally to minimize the step number.

  Note that the agents can only lift one box at a time. Each lifting agent can be used only once in each step! You can combine multiple agents to lift one box like "box[3.0V]":"agent[1.5W], agent[2.5W]"! Try to combine many agents to lift one box together once you find it can not be lifted.
  [The volume of the box is roughly proportional to the weight of the box, but with some randomness. Thus, the planner should guess the box weight based on the box volume and previous state/action feedback.]

  Specify your action plan in the JSON format: {{"box[1.7V]":"agent[1.5W]", "box[3.0V]":"agent[1.5W], agent[2.5W], agent[5.5W]"}}. Include a box only if it has lifting agents to lift it next.

  '''

if not os.path.exists(base_path):
    raise ValueError(f'base_path {base_path} does not exist')

dir_path = os.path.join(base_path, f'prompt_optimization_train_result')
if not os.path.exists(dir_path):
    os.makedirs(dir_path)
save_path = os.path.join(dir_path, f'train_result_dir_trial_{experiment_trial_num}_PromptLLM_{model_name_promptLLM}_TestLLM_{model_name_testLLM}_PromptMethod_{prompt_method}')

tree = Tree(base_path=save_path, Testing_path = Testing_path, n_children=n_children, n_selected=n_selected,
            input_error_prompt_token_limit = input_error_prompt_token_limit, model_name_promptLLM = model_name_promptLLM, model_name_testLLM = model_name_testLLM, prompt_method = prompt_method, with_score_model = with_score_model)

# Check if the directory exists
if os.path.exists(save_path):
    # Remove the directory and its contents
    shutil.rmtree(save_path)

# Create the directory
os.makedirs(save_path)

path_each_prompt = os.path.join(save_path, f'prompt_{tree.node_index}')
if not os.path.exists(path_each_prompt):
    os.makedirs(path_each_prompt)
score, feedback_to_promptLLM_list, success_failure_list, error_string_list = tree.score_prompt(Starting_prompt_task_explain, path_each_prompt)
tree.root = Node(Starting_prompt_task_explain, score, 1, tree.node_index, parent=None,
                 feedback_to_promptLLM_list=feedback_to_promptLLM_list, success_failure_list=success_failure_list, error_string_list=error_string_list)

tree.evolve_tree(tree.root, min_level=min_level)
tree.display_tree(tree.root)  # Display the tree for visualization

# Save the tree to a file
tree.save_tree(tree.root, save_path)
tree.save_display_tree(tree.root)
