import numpy as np
from perception_utils import parse_query_obj
from plan_utils import get_empty_cost_map, set_voxel_by_radius, cm2index

# Query: a point 10cm in front of [10, 15, 60].
cost_map = get_empty_cost_map()
# 10cm in front of so we add to x-axis
x = 10 + cm2index(10, 'x')
y = 15
z = 60
cost_map[x, y, z] = 1
ret_val = cost_map

# Query: a point on the right side of the table.
cost_map = get_empty_cost_map()
table = parse_query_obj('table')
(min_x, min_y, min_z), (max_x, max_y, max_z) = table.aabb
center_x, center_y, center_z = table.position
# right side so y = max_y
x = center_x
y = max_y
z = center_z
cost_map[x, y, z] = 1
ret_val = cost_map

# Query: a point 20cm on top of the container.
cost_map = get_empty_cost_map()
container = parse_query_obj('container')
(min_x, min_y, min_z), (max_x, max_y, max_z) = container.aabb
center_x, center_y, center_z = container.position
# 20cm on top of so we add to z-axis
x = center_x
y = center_y
z = max_z + cm2index(20, 'z')
cost_map[x, y, z] = 1
ret_val = cost_map

# Query: a point 1cm to the left of the brown block.
cost_map = get_empty_cost_map()
brown_block = parse_query_obj('brown block')
(min_x, min_y, min_z), (max_x, max_y, max_z) = brown_block.aabb
center_x, center_y, center_z = brown_block.position
# 1cm to the left of so we subtract from y-axis
x = center_x
y = min_y - cm2index(1, 'y')
z = center_z
cost_map[x, y, z] = 1
ret_val = cost_map

# Query: anywhere within 20cm of the right most block.
cost_map = get_empty_cost_map()
right_most_block = parse_query_obj('the right most block')
set_voxel_by_radius(cost_map, right_most_block.position, radius_cm=20, value=1)

# Query: a point on the back side of the table.
cost_map = get_empty_cost_map()
table = parse_query_obj('table')
(min_x, min_y, min_z), (max_x, max_y, max_z) = table.aabb
center_x, center_y, center_z = table.position
# back side so x = min_x
x = min_x
y = center_y
z = center_z
cost_map[x, y, z] = 1
ret_val = cost_map

# Query: a point on the front right corner of the table.
cost_map = get_empty_cost_map()
table = parse_query_obj('table')
(min_x, min_y, min_z), (max_x, max_y, max_z) = table.aabb
center_x, center_y, center_z = table.position
# front right corner so x = max_x and y = max_y
x = max_x
y = max_y
z = center_z
cost_map[x, y, z] = 1
ret_val = cost_map

# Query: a point 30cm into the topmost drawer handle.
cost_map = get_empty_cost_map()
top_handle = parse_query_obj('topmost drawer handle')
# negative normal because we are moving into the handle.
moving_dir = -top_handle.normal
cost_xyz = top_handle.position + cm2index(30, moving_dir)
cost_map[cost_xyz[0], cost_xyz[1], cost_xyz[2]] = 1
ret_val = cost_map

# Query: a point 5cm above the blue block.
cost_map = get_empty_cost_map()
blue_block = parse_query_obj('blue block')
(min_x, min_y, min_z), (max_x, max_y, max_z) = blue_block.aabb
center_x, center_y, center_z = blue_block.position
# 5cm above so we add to z-axis
x = center_x
y = center_y
z = max_z + cm2index(5, 'z')
cost_map[x, y, z] = 1
ret_val = cost_map

# Query: a point 4cm to the left of and 10cm on top of the tray that contains the lemon.
cost_map = get_empty_cost_map()
tray_with_lemon = parse_query_obj('tray that contains the lemon')
(min_x, min_y, min_z), (max_x, max_y, max_z) = tray_with_lemon.aabb
center_x, center_y, center_z = tray_with_lemon.position
# 4cm to the left of so we subtract from y-axis, and 10cm on top of so we add to z-axis
x = center_x
y = min_y - cm2index(4, 'y')
z = max_z + cm2index(10, 'z')
cost_map[x, y, z] = 1
ret_val = cost_map

# Query: a point 10cm to the right of [45 49 66], and 5cm above it.
cost_map = get_empty_cost_map()
# 10cm to the right of so we add to y-axis, and 5cm above it so we add to z-axis
x = 45
y = 49 + cm2index(10, 'y')
z = 66 + cm2index(5, 'z')
cost_map[x, y, z] = 1
ret_val = cost_map

# Query: the blue circle.
cost_map = get_empty_cost_map()
blue_circle = parse_query_obj('blue circle')
cost_map = blue_circle.occupancy_map
ret_val = cost_map

# Query: a point at the center of the blue circle.
cost_map = get_empty_cost_map()
blue_circle = parse_query_obj('blue circle')
x, y, z = blue_block.position
cost_map[x, y, z] = 1
ret_val = cost_map

# Query: a point 10cm above and 5cm to the left of the yellow bowl.
cost_map = get_empty_cost_map()
yellow_bowl = parse_query_obj('yellow bowl')
(min_x, min_y, min_z), (max_x, max_y, max_z) = yellow_bowl.aabb
center_x, center_y, center_z = yellow_bowl.position
# 10cm above so we add to z-axis, and 5cm to the left of so we subtract from y-axis
x = center_x
y = min_y - cm2index(5, 'y')
z = max_z + cm2index(10, 'z')
cost_map[x, y, z] = 1
ret_val = cost_map

# Query: a point on the side of the block with its back to the green_target.
cost_map = get_empty_cost_map()
block = parse_query_obj('block')
block_x, block_y, block_z = block.position
green_target = parse_query_obj('green_target')
green_target_x, green_target_y, green_target_z = green_target.position
# 5cm on the side of the block with its back to the green_target
direction = np.array([green_target_x - block_x, green_target_y - block_y, 0])  
cost_xyz = block.position - cm2index(15, direction)
cost_map[cost_xyz[0], cost_xyz[1], cost_xyz[2]] = 1
ret_val = cost_map

# Query: a point 30cm away from the bottom_handle.
cost_map = get_empty_cost_map()
bottom_handle = parse_query_obj('bottom_handle')
bottom_x, bottom_y, bottom_z = bottom_handle.position
x = bottom_x
y = bottom_y + cm2index(30, 'y') # y axis: a point 30cm away from bottom_handle
z = bottom_z
cost_map[x, y, z] = 1
ret_val = cost_map

# Query: a point at the center of jar_closer_to_banana.
cost_map = get_empty_cost_map()
bottom_handle = parse_query_obj('jar_closer_to_banana')
x, y, z = bottom_handle.position
cost_map[x, y, z] = 1
ret_val = cost_map