from cozmo_fsm import *

"""
    Drive forward 25 mm and measure actual forward travel and sideways drift.
    Display as a scatter plot.
"""

import matplotlib.pyplot as plt

class PlotDrive(StateMachineProgram):

    class Setup(StateNode):
        def start(self,event=None):
            super().start()
            self.parent.fig = plt.figure()
            self.parent.ax = ax = self.parent.fig.add_subplot(111)
            ax.set_xlim(-5,5)
            ax.set_ylim(23,29)
            ax.plot([-5,5],[25,25],'g--')
            ax.plot([0,0],[23,29],'r--')
            ax.set_xlabel('Drift')
            ax.set_ylabel('Travel Distance')
            self.parent.old_pos = self.robot.pose.position


    class AddPoint(StateNode):
        def start(self,event=None):
            super().start()
            fwd = self.robot.pose.position.x - self.parent.old_pos.x
            drift = self.robot.pose.position.y - self.parent.old_pos.y
            print('fwd=',fwd,'drift=',drift)
            self.parent.ax.plot(drift,fwd,'bo',linestyle='None')
            self.parent.old_pos = self.robot.pose.position

    class ShowPlot(StateNode):
        def start(self,event=None):
            super().start()
            plt.show()

    def setup(self):
        """
            self.Setup() =N=> loop
    
            loop: Iterate(range(5))
            loop =D=> Forward(25) =C=> self.AddPoint() =Next=> loop
            loop =C=> show
    
            show: self.ShowPlot()
    
        """
        
        # Code generated by genfsm on Mon Feb 18 05:50:48 2019:
        
        setup1 = self.Setup() .set_name("setup1") .set_parent(self)
        loop = Iterate(range(5)) .set_name("loop") .set_parent(self)
        forward1 = Forward(25) .set_name("forward1") .set_parent(self)
        addpoint1 = self.AddPoint() .set_name("addpoint1") .set_parent(self)
        show = self.ShowPlot() .set_name("show") .set_parent(self)
        
        nulltrans1 = NullTrans() .set_name("nulltrans1")
        nulltrans1 .add_sources(setup1) .add_destinations(loop)
        
        datatrans1 = DataTrans() .set_name("datatrans1")
        datatrans1 .add_sources(loop) .add_destinations(forward1)
        
        completiontrans1 = CompletionTrans() .set_name("completiontrans1")
        completiontrans1 .add_sources(forward1) .add_destinations(addpoint1)
        
        nexttrans1 = NextTrans() .set_name("nexttrans1")
        nexttrans1 .add_sources(addpoint1) .add_destinations(loop)
        
        completiontrans2 = CompletionTrans() .set_name("completiontrans2")
        completiontrans2 .add_sources(loop) .add_destinations(show)
        
        return self