from cozmo_fsm import *

class CircleWatcher(StateMachineProgram):

    def user_image(self,image,gray):
        ret, thresh = cv2.threshold(gray, 50, 255, 0)
        stuff, 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 c in self.indices:
            if c[1] < 500:  # too small
                break
            if c[2] < 0.8:  # not solid enough
                continue
            i = c[0]
            cnt = scale * self.contours[i]
            ellipse = cv2.fitEllipse(cnt)
            cv2.drawContours(annotated_image, [cnt], 0, (255, 255, 255), 3)
            cv2.ellipse(annotated_image, ellipse, (255,0,255), 1)
        return annotated_image