from utils import convert_2d_to_3d
from inference_utils import mllm_inference
from perception_utils import parse_constr_coords, get_mask_edge_coords, transform_to_label, transform_to_coords, pos_constr_refine

#Query: Get the extracted constraints, and refine the constraint for the red mug’s opening
#positional constraint refinement
constr_coords = parse_constr_coords("red_mug")["extracted"]
edge_coords = get_mask_edge_coords(constr_coords)
refined_pos_constr = pos_constr_refine(edge_coords)
ret_val = [refined_pos_constr]

#Query: Get the extracted constraints, and refine the constraints for both the stick's thick end and the evaporating dish's opening.
#geometry constraint refinement
constr_coords = parse_constr_coords("stick")["extracted"]
label = transform_to_label(constr_coords)
refined_label = mllm_inference("Refine the constraints for the stick's thick end.", type="geometry_refine", params=[constr_coords[-1], label])
refined_coords = transform_to_coords(refined_label)
refined_geom_constr = convert_2d_to_3d(refined_coords)
#positional constraint refinement
constr_coords = parse_constr_coords("evaporating_dish")["extracted"]
edge_coords = get_mask_edge_coords(constr_coords)
refined_pos_constr = pos_constr_refine(edge_coords)
ret_val = [refined_geom_constr, refined_pos_constr]

#Query: Get the extracted constraints, and refine the constraint for the pen's cap.
#geometry constraint refinement
constr_coords = parse_constr_coords("pen")["extracted"]
label = transform_to_label(constr_coords)
refined_label = mllm_inference("Refine the constraint for the pen's cap.", type="geometry_refine", params=[constr_coords[-1], label])
refined_coords = transform_to_coords(refined_label)
refined_geom_constr = convert_2d_to_3d(refined_coords)
ret_val = [refined_geom_constr]