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)

    $setup{
        self.Setup() =N=> loop

        loop: Iterate(range(5))
        loop =D=> Forward(25) =T(3)=> self.AddPoint() =Next=> loop
        loop =C=> show

        show: self.ShowPlot()

    }