from inference_utils import mllm_inference, constraint_refinement
from perception_utils import get_2d_coords, map_constraints, merge_constraints_state, get_3d_state, register_save_constraints

#Query: Extract constraints for bottle’s opening and electronic balance's center, and refine the constraint for the bottle's opening.
objects = ["bottle", "electronic_balance"]
extracted_constr = mllm_inference("Extract constraints for bottle’s opening and electronic balance's center.", type="extract", objects=objects)
constr_coords = get_2d_coords(extracted_constr)
assert map_constraints(constr_coords, objects), "constraints don't match the objects."
refined_constr = constraint_refinement("Get the extracted constraints, and refine the constraint for the bottle’s opening")
constr_state = get_3d_state(constr_coords)
constr_state = merge_constraints_state(constr_state, refined_constr, objects=["bottle"], type="refined")
register_save_constraints(constr_state)

#Query: Extract constraints for pen’s center and pen holder's center, and refine the constraints for both the pen's thick end and the pen holder's opening.
objects = ["pen", "pen_holder"]
extracted_constr = mllm_inference("Extract constraints for pen’s center and pen holder's center.", type="extract", objects=objects)
constr_coords = get_2d_coords(extracted_constr)
assert map_constraints(constr_coords, objects), "constraints don't match the objects."
refined_constr = constraint_refinement("Get the extracted constraints, and refine the constraints for both the pen's thick end and the pen holder's opening")
constr_state = get_3d_state(constr_coords)
constr_state = merge_constraints_state(constr_state, refined_constr, objects=["pen", "pen_holder"], type="refined")
register_save_constraints(constr_state)

#Query: Extract constraints for trash.
objects = ["trash"]
extracted_constr = mllm_inference("Extract constraints for trash.", type="extract", objects=objects)
constr_coords = get_2d_coords(extracted_constr)
assert map_constraints(constr_coords, objects), "constraints don't match the objects."
constr_state = get_3d_state(constr_coords)
register_save_constraints(constr_state)