import cv2
import robot

def find_aruco(arlo):
    image = arlo.take_photo()

    aruco_dict = cv2.aruco.Dictionary_get(cv2.aruco.DICT_6X6_250)
    aruco_params = cv2.aruco.DetectorParameters_create()
    corners, ids, _ = cv2.aruco.detectMarkers(
        image,
        aruco_dict,
        parameters=aruco_params
    )

    for bounding, n in zip(corners, ids):
        bounding_box = [(int(x), int(y)) for x, y in bounding.reshape((4,2))]
        print(bounding_box)

        cv2.line(image, bounding_box[0], bounding_box[1], (0,255,0), 2)
        cv2.line(image, bounding_box[1], bounding_box[2], (0,255,0), 2)
        cv2.line(image, bounding_box[2], bounding_box[3], (0,255,0), 2)
        cv2.line(image, bounding_box[3], bounding_box[0], (0,255,0), 2)

        cv2.putText(
            image,
            str(n[0]),
            (bounding_box[0][0] - 10, bounding_box[0][1] - 10),
            cv2.FONT_HERSHEY_COMPLEX,
            0.8,
            (0,255,0),
            2
        )

    return image

def main():
    arlo = robot.Robot()
    image = find_aruco(arlo)
    cv2.imwrite("test.png", image)

if __name__ == "__main__":
    main()