import numpy as np
from env_utils import reset_to_default_pose
from perception_utils import parse_query_obj
from plan_utils import get_cost_map, get_gripper_map

# Query: move ee forward for 10cm then open gripper.
movable = parse_query_obj('gripper')
cost_map = get_cost_map(f'a point 10cm in front of {movable.position}')
gripper_map = get_gripper_map(f'close everywhere except 3cm around a point 10cm in front of {movable.position}')

# Query: drop the toy inside container.
movable = parse_query_obj('gripper')
cost_map = get_cost_map('a point 15cm above the container')
gripper_map = get_gripper_map('close everywhere but open when on top of the container')

# Query: push close the topmost drawer.
movable = parse_query_obj('topmost drawer handle')
cost_map = get_cost_map('a point 30cm into the topmost drawer handle')
gripper_map = get_gripper_map('close everywhere')

# Query: push the second to the left block along the red line.
movable = parse_query_obj('second to the left block')
cost_map = get_cost_map('the red line')
gripper_map = get_gripper_map('close everywhere')

# Query: grasp the blue block from the table.
movable = parse_query_obj('gripper')
cost_map = get_cost_map('a point at the center of blue block')
gripper_map = get_gripper_map('open everywhere except 1cm around the blue block')

# Query: move to the left of the brown block.
movable = parse_query_obj('gripper')
cost_map = get_cost_map('a point 10cm to the left of the brown block')
gripper_map = get_gripper_map('open everywhere')

# Query: move to the top of the tray that contains the lemon.
movable = parse_query_obj('gripper')
cost_map = get_cost_map('a point 10cm above the tray that contains the lemon')
gripper_map = get_gripper_map('oprn everywhere')

# Query: close drawer by 5cm.
movable = parse_query_obj('drawer handle')
cost_map = get_cost_map('a point 5cm into the drawer handle')
gripper_map = get_gripper_map('close everywhere')

# Query: move to 5cm on top of the soda can then open gripper.
movable = parse_query_obj('gripper')
cost_map = get_cost_map('a point 5cm above the soda can')
gripper_map = get_gripper_map('close everywhere except 3cm around the 5cm on top of the soda')

# Query: wipe the red dot.
movable = parse_query_obj('gripper')
cost_map = get_cost_map('the red dot')
gripper_map = get_gripper_map('close everywhere')

# Query: grasp the mug from the shelf.
movable = parse_query_obj('gripper')
cost_map = get_cost_map('a point at the center of the mug handle')
gripper_map = get_gripper_map('open everywhere except 3cm around the mug handle')

# Query: open gripper.
movable = parse_query_obj('gripper')
gripper_map = get_gripper_map('open everywhere')