from cozmo_fsm import *

class FindChips(StateNode):
    def start(self, event=None):
        super().start(event)
        thresh1 = 130  # was 140
        minArea = 200
        chips = []
        raw_img = self.robot.world.latest_image.raw_image
        image = numpy.array(raw_img).astype(numpy.uint8)
        csvImg = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
        lower_red1 = np.array([0, thresh1, thresh1])
        upper_red1 = np.array([10, 255, 255])
        image1 = cv2.inRange(csvImg, lower_red1, upper_red1)
        lower_red2 = np.array([170, thresh1, thresh1])
        upper_red2 = np.array([180,255,255])
        image2 = cv2.inRange(csvImg, lower_red2, upper_red2)
        mask = cv2.bitwise_or(image1, image2)
        cv2.namedWindow('mask')
        cv2.imshow('mask',mask)
        stuff, contours, hierarchy = \
            cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        areas = [(i, cv2.contourArea(contours[i])) for i in range(len(contours))]
        areas.sort(key=lambda x: x[1])
        areas.reverse()
        for area_entry in areas:
            if area_entry[1] < minArea:
                break
            index = area_entry[0]
            contour = 2 * contours[index]
            if len(contour) < 5:
                continue
            ellipse = cv2.fitEllipse(contour)
            area = cv2.contourArea(contour)
            earea = pi*ellipse[1][1]*ellipse[1][0]/4.
            solidity = float(area)/earea if earea > 0 else 0
            M = cv2.moments(contour)
            cx = int(M['m10']/M['m00'])
            cy = int(M['m01']/M['m00'])
            if solidity > 0.75:
                if ( not (cy > 375 or cy < 150)
                        and not (cx < 50 and cy > 225 and cy < 275)):
                    (x,y),radius = cv2.minEnclosingCircle(contour)
                    dist = (19 * self.robot.camera.config.focal_length.x) / radius
                    if (cx == 320.):
                        slope = 1000000
                    else:
                        slope = (cy - 1200.) / (cx - 320.)
                    angleFactor = 8.5 / 0.1
                    angle = 1/slope * angleFactor * pi/180
                    rX = dist * sin(angle)
                    rY = dist * cos(angle)
                    chips.append((rX, rY, radius, earea))
        self.post_data(chips)

class Dumper(StateNode):
    def start(self, event=None):
        super().start(event)
        print('Chips: ', end='')
        for coords in event.data:
            print('[%.1f %.1f] R=%.1f A=%.1f   ' % coords, end='')
        print()
        

class Chipper(StateMachineProgram):
  def setup(self):
      """
        ColorImageEnabled(True) =C=> find: FindChips() =D=> Dumper() =T(1)=> find
      """
      
      # Code generated by genfsm on Fri Mar  2 07:59:49 2018:
      
      colorimageenabled1 = ColorImageEnabled(True) .set_name("colorimageenabled1") .set_parent(self)
      find = FindChips() .set_name("find") .set_parent(self)
      dumper1 = Dumper() .set_name("dumper1") .set_parent(self)
      
      completiontrans1 = CompletionTrans() .set_name("completiontrans1")
      completiontrans1 .add_sources(colorimageenabled1) .add_destinations(find)
      
      datatrans1 = DataTrans() .set_name("datatrans1")
      datatrans1 .add_sources(find) .add_destinations(dumper1)
      
      timertrans1 = TimerTrans(1) .set_name("timertrans1")
      timertrans1 .add_sources(dumper1) .add_destinations(find)
      
      return self