from aim_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):
    def __init__(self):
        super().__init__(launch_cam_viewer=False,
		         launch_worldmap_viewer=False)

    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 (deg.)')
            ax.set_ylabel('Travel Distance')
            self.parent.old_pose = self.robot.pose


    class AddPoint(StateNode):
        def start(self,event=None):
            super().start()
            fwd = self.robot.pose.x - self.parent.old_pose.x
            drift = wrap_angle(self.robot.pose.theta - self.parent.old_pose.theta)*180/pi
            print('fwd=',fwd,'drift=',drift)
            self.parent.ax.plot(drift,fwd,'bo',linestyle='None')
            self.parent.old_pose = self.robot.pose

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

    def setup(self):
        #         self.Setup() =N=> loop
        # 
        #         loop: Iterate(range(5))
        #         loop =D=> Forward(25) =T(3)=> self.AddPoint() =Next=> loop
        #         loop =C=> show
        # 
        #         show: self.ShowPlot()
        # 
        
        # Code generated by genfsm on Mon Feb 17 02:53:54 2025:
        
        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)
        
        timertrans1 = TimerTrans(3) .set_name("timertrans1")
        timertrans1 .add_sources(forward1) .add_destinations(addpoint1)
        
        nexttrans1 = NextTrans() .set_name("nexttrans1")
        nexttrans1 .add_sources(addpoint1) .add_destinations(loop)
        
        completiontrans1 = CompletionTrans() .set_name("completiontrans1")
        completiontrans1 .add_sources(loop) .add_destinations(show)
        
        return self