from LLM import *
from prompt_env9 import *
from env9_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 env9_func import *
from argparse import ArgumentParser

import pdb
import sys
import re
import wandb
import warnings
import yaml
import time
from dotenv import load_dotenv
from tasks import load_task
from llm import load_llm

parser = ArgumentParser()
## Extra parameters for alfworld
parser.add_argument("--cfg-path", required=True, help="path to configuration file.")
parser.add_argument("--tasks", required=True, type=str, nargs='+', help="specify the tasks")
parser.add_argument("--wandb", action="store_true", help="specify whether the wandb board is needed")
parser.add_argument("--log_path", required=False, default='', help="specify the place to store the resuls")
parser.add_argument("--project_name", required=False, default='', help="specify the project name for wandb")
parser.add_argument("--baseline_dir", required=False, default='',
                    help="specify the baseline loggings for wandb baseline comparison visualization")
parser.add_argument('--max_num_steps', '--max_num_steps', type=int, default=30) # specify the maximum number of steps used to finish the problems

## Original parameters for PROMST
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('-base_path', '--base_path', default='../BoxNet1/') # 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
base_path = args.base_path
n_children = args.n_children
n_selected = args.n_selected

if args.tasks[0] == 'alfworld':
    print('#'*20)
    print(f'Task: {args.tasks[0]}')
    Starting_prompt_task_explain = f'''
    Your task is to interact with a virtual household simulator to accomplish a specific task. With each interaction, you will receive an observation and current valid actions.
    Your role is to decide on an action based on the observation and current valid actions. Please ensure that any objects ('{{obj}}') and receptacles ('{{recep}}') you mention in your response are present in the observation provided.
    Ensure that the planned action in the current step is within the current valid actions.
    Example objects are like a cellphone 3, a newspaper 2, a statue 1, and a television 1.
    Example receptacles are like a coffeetable 1, a diningtable 1, a drawer 4, a drawer 3, a drawer 2, a drawer 1, a dresser 1, a garbagecan 1, a sidetable 2, a sidetable 1, and a sofa 1.
    Example actions are like [go to dresser 1, take statue 1 from dresser 1, heat apple 1 with microwave 1, open cabinet 2]
    Do not repeat the actions all the time! Learn from the previous action/observation history.
    
    Here are the available actions you can take:
    - 'take {{obj}} from {{recep}}'
    - 'put {{obj}} in/on {{recep}}'
    - 'open {{recep}}'
    - 'close {{recep}}'
    - 'toggle {{obj}}/{{recep}}'
    - 'clean {{obj}} using {{recep}}'
    - 'cool {{obj}} using {{recep}}'
    - 'heat {{obj}} using {{recep}}'
    - 'inventory'
    - 'examine {{recep}}/{{obj}}'
    - 'go to {{recep}}'
    
              '''
elif args.tasks[0] == 'scienceworld':
    print('#'*20)
    print(f'Task: {args.tasks[0]}')
    Starting_prompt_task_explain = f'''
    You are an agent in a virtual science school environment, tasked to interact with various elements.
    
    Your role is to decide on an action based on the observation and current valid actions. Please ensure that any objects ('{{OBJ}}') and locations ('{{LOC}}') you mention in your response are present in the observation provided.
    Ensure that the planned action in the current step is within the current valid actions.
    Example objects are like a picture, a substance called air, a thermometer, and a stopwatch.
    Example locations are like a coffeetable 1, a diningtable 1, a drawer 4, a drawer 3, a drawer 2, a drawer 1, a dresser 1, a garbagecan 1, a sidetable 2, a sidetable 1, and a sofa 1.
    Example actions are like [go to dresser 1, take statue 1 from dresser 1, heat apple 1 with microwave 1, open cabinet 2]
    Do not repeat the actions all the time! Learn from the previous action/observation history.
    
    Here are the commands you can use:
    - **Manipulation**: 
      - `open {{OBJ}}` / `close {{OBJ}}`: Interact with a container.
      - `pick up {{OBJ}}`: Add an object to your inventory.
      - `put down {{OBJ}}`: Remove an object from your inventory.
      - `move {{OBJ}} to {{OBJ}}`: Transfer an object.
      - `pour {{OBJ}} into {{OBJ}}`: Pour a substance.
      - `dunk {{OBJ}} into {{OBJ}}`: Immerse a container in a liquid.
      - `mix {{OBJ}}`: Chemically combine contents.
    
    - **Inspection**:
      - `look around`: Survey your surroundings.
      - `look at {{OBJ}}`: Examine an object closely.
      - `look in {{OBJ}}`: Peek inside a container.
      - `read {{OBJ}}`: Review written content.
    
    - **Device Operations**:
      - `activate {{OBJ}}` / `deactivate {{OBJ}}`: Toggle a device.
      - `use {{OBJ}} [on {{OBJ}}]`: Utilize a device or item.
    
    - **Movement**:
      - `go to {{LOC}}`: Relocate.
    
    - **Miscellaneous**:
      - `eat {{OBJ}}`: Consume an edible item.
      - `flush {{OBJ}}`: Activate a flushing mechanism.
      - `focus on {{OBJ}}`: Direct attention to a particular object.
      - `wait [DURATION]`: Pause for a specified period.
    
    - **Information**:
      - `task`: Recap your current objective.
      - `inventory`: Display items you're carrying.
    
    Where:
    - `{{OBJ}}`: Object
    - `{{LOC}}`: Location
    - `[DURATION]`: Specified time
    
              '''
elif args.tasks[0] == 'webarena':
    print('#'*20)
    print(f'Task: {args.tasks[0]}')
    Starting_prompt_task_explain = f'''
    Here's the information you'll have:
    The user's objective: This is the task you're trying to complete.
    The current web page's accessibility tree: This is a simplified representation of the windowed webpage, providing key information.
    The current web page's URL: This is the page you're currently navigating.
    The open tabs: These are the tabs you have open.
    
    The useful websites and corresponding URL you can navigate:
    'reddit': "http://reddit.com"
    'online shop': "http://onestopmarket.com"
    'e-commerce platform': "http://luma.com/admin"
    'gitlab': "http://gitlab.com"
    'wikipedia': "http://wikipedia.org"
    'map': "http://openstreetmap.org"
    
    Your role is to decide on an action based on the observation and current valid actions.
    Ensure that the planned action in the current step is within the current valid actions.
    
    The actions you can perform fall into several categories:
    
    Page Operation Actions:
    `click [id]`: This action clicks on an element with a specific id on the webpage.
    `type [id] [content] [press_enter_after=0|1]`: Use this to type the content into the field with id. By default, the "Enter" key is pressed after typing unless press_enter_after is set to 0.
    `hover [id]`: Hover over an element with id.
    `press [key_comb]`:  Simulates the pressing of a key combination on the keyboard (e.g., Ctrl+v).
    `scroll [direction=down|up]`: Scroll the page up or down.
    
    Tab Management Actions:
    `new_tab`: Open a new, empty browser tab.
    `tab_focus [tab_index]`: Switch the browser's focus to a specific tab using its index.
    `close_tab`: Close the currently active tab.
    
    URL Navigation Actions:
    `goto [url]`: Navigate to a specific URL.
    `go_back`: Navigate to the previously viewed page.
    `go_forward`: Navigate to the next page (if a previous 'go_back' action was performed).
    
    Completion Action:
    `stop [answer]`: Apply this action when you believe the task is complete. If it is a operation-type task, use `stop [Done]` when finished. If the objective is to give a text-based answer, provide the answer in the bracket. 
    
    To be successful, it is very important to follow the following rules:
    1. You should only issue an action that is valid given the current observation
    2. You should only issue one action at a time.
    3. Generate the action in the correct format and always put the action inside a pair of @. Such as, @click [1234]@.
    4. Complete the task by interacting with the starting page, and avoid using 'goto' actions casually.
    5. Reasonable inputs will return accurate observations, so do not repeat the same action when unnecessary.
    
    The following is one task executing example:
    You should perform actions to accomplish the goal: Add a white desk to my wish list
    Observation:
    WINDOWED PAGE:{{
    Tab 0 (current): One Stop Market
    
    [1] RootWebArea 'One Stop Market' focused: True
        [1254] link 'My Wish List'
        [2427] StaticText 'Search'
        [1585] combobox 'Search' autocomplete: both hasPopup: listbox required: False expanded: False
        [2430] link 'Advanced Search'
        [1588] button 'Search' disabled: True
    }}
    URL: http://onestopmarket.com
    
    Action: @type [1585] [white desk] [press_enter_after=1]@
    Observation:
    WINDOWED PAGE:{{
    Tab 0 (current): Search results for: 'white desk'
    
    [2635] RootWebArea "Search results for: 'white desk'" focused: True
        [3869] link 'My Wish List'
        [4827] StaticText 'Search'
        [4072] combobox 'Search' autocomplete: both hasPopup: listbox required: False expanded: False
            [5027] StaticText 'white desk'
        [4830] link 'Advanced Search'
        [4075] button 'Search' disabled: True
        [3729] main ''
            [3842] heading "Search results for: 'white desk'"
            [3907] StaticText 'Items 1-12 of 38823'
            [4781] link 'Image'
                [4833] img 'Image'
            [4783] link 'Image'
                [4849] img 'Image'
    }}
    URL: http://onestopmarket.com/catalogsearch/result/?q=white+desk
    
    Action: @scroll [down]@
    Observation:
    WINDOWED PAGE:{{
    Tab 0 (current): Search results for: 'white desk'
    
    [2635] RootWebArea "Search results for: 'white desk'" focused: True
        [3729] main ''
            [4781] link 'Image'
                [4833] img 'Image'
            [4838] link 'Techni Mobili Rolling Writing Height Adjustable Desktop and Moveable Shelf,White Desk'
            [5034] StaticText '$123.99'
            [5038] button 'Add to Cart'
            [4783] link 'Image'
                [4849] img 'Image'
            [4853] link 'VIVO Electric Motor Height Adjustable Corner 3 Leg Standing Desk Frame, Frame Only, Sit Stand Ergonomic L Frame, White, DESK-V130EW'
            [5043] StaticText '$499.99'
            [5047] button 'Add to Cart'
            [4785] link 'Image'
                [4862] img 'Image'
            [4787] link 'Image'
                [4875] img 'Image'
    }}
    URL: http://onestopmarket.com/catalogsearch/result/?q=white+desk
    
    Action: @scroll [down]@
    Observation:
    WINDOWED PAGE:{{
    Tab 0 (current): My Wish List
    
    [5381] RootWebArea 'My Wish List' focused: True
        [6685] link 'My Wish List 1 items'
        [6470] main ''
            [6605] alert '' live: assertive atomic: True relevant: additions text
                [6777] StaticText 'Techni Mobili Rolling Writing Height Adjustable Desktop and Moveable Shelf,White Desk has been added to your Wish List. Click '
                [6778] link 'here'
                [6779] StaticText ' to continue shopping.'
    }}
    URL: http://onestopmarket.com/wishlist/index/index/wishlist_id/5/
    
    Action: @stop [Done]@
    Observation:
    Well Done!
    URL: http://onestopmarket.com/wishlist/index/index/wishlist_id/5/
    
    
    '''
else:
    print('#'*20)
    print(f'No environment specified!')
    print(f'Task: {args.tasks[0]}')

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'round1_{model_name_promptLLM}_{model_name_testLLM}_wo_score_model_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(args, base_path=save_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)
