from aim_fsm import *

class CircleWatcher(StateMachineProgram):
    def __init__(self):
        super().__init__(aruco=False,
                         launch_worldmap_viewer=False,
                         force_annotation=True,
                         annotate_sdk=False,
                         speech=False)

    def user_image(self, image, gray):
        ret, thresh = cv2.threshold(gray, 50, 255, cv2.THRESH_BINARY)
        contours, hierarchy = \
            cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        # Each index entry: (index, area, solidity)
        indices = [(i,
                    cv2.contourArea(contours[i]),
                    float(cv2.contourArea(contours[i])) /
                      max(1,cv2.contourArea(cv2.convexHull(contours[i]))))
                   for i in range(len(contours))]
        indices.sort(key=lambda x: x[1])
        indices.reverse()
        self.indices = indices
        self.contours = contours

    def user_annotate(self, annotated_image):
        scale = self.annotated_scale_factor
        for contour_index in self.indices:
            if contour_index[1] < 500:  # too small
                break
            elif contour_index[1] > 20000: # too large
                continue
            elif contour_index[2] < 0.7:  # not solid enough
                continue
            # print(contour_index)
            i = contour_index[0]
            contour = scale * self.contours[i]
            ellipse = cv2.fitEllipse(contour)
            cv2.drawContours(annotated_image, [contour], 0, (255, 255, 255), 3)
            cv2.ellipse(annotated_image, ellipse, (255,0,255), 5)
        return annotated_image