from cozmo_fsm import *

class Lab6(StateMachineProgram):
    def start(self):
        start = RRTNode(x=-100, y=-100)
        goal = RRTNode(x=200, y=200, q=math.nan)
        obst1 = Circle(center=transform.point(100,-20), radius=20)
        obst2 = Circle(center=transform.point(100,100), radius=30)
        obst3 = Circle(center=transform.point(-30,100), radius=40)
        self.rrt = RRT(robot, step_size=10, obstacles=[obst1,obst2,obst3],auto_obstacles=False)

        self.pv = PathViewer(self.robot,self.rrt)
        self.pv.start()

        (treeA, treeB, path) = self.rrt.plan_path(start,goal)
        print(len(treeA)+len(treeB),'nodes')

        # display the path in yellow, and start/end nodes in white
        self.pv.add_tree(path, (1,1,0,0.5))
        self.pv.add_tree([start], (1,1,1))
        self.pv.add_tree([goal], (1,1,1))