Make the Spline UI widgets more widgety

Change-Id: I89d9ad7b795cef716f6da55414175acbf49c1855
Signed-off-by: Ravago Jones <ravagojones@gmail.com>
diff --git a/frc971/control_loops/python/graph.py b/frc971/control_loops/python/graph.py
index 3063bd1..1b83713 100644
--- a/frc971/control_loops/python/graph.py
+++ b/frc971/control_loops/python/graph.py
@@ -1,209 +1,55 @@
-from constants import *
-import cairo
-from color import Color, palette
+import gi
+gi.require_version('Gtk', '3.0')
+from gi.repository import Gtk
+import numpy as np
 from points import Points
-from drawing_constants import *
 from libspline import Spline, DistanceSpline, Trajectory
 
-AXIS_MARGIN_SPACE = 40
+from matplotlib.backends.backend_gtk3agg import (FigureCanvasGTK3Agg as
+                                                 FigureCanvas)
+from matplotlib.figure import Figure
 
 
-class Graph():  # (TODO): Remove Computer Calculation
-    def __init__(self, cr, mypoints):
-        # Background Box
-        set_color(cr, palette["WHITE"])
-        cr.rectangle(-1.0 * SCREEN_SIZE, -0.5 * SCREEN_SIZE, SCREEN_SIZE,
-                     SCREEN_SIZE * 0.6)
-        cr.fill()
+class Graph(Gtk.Bin):
+    def __init__(self):
+        super(Graph, self).__init__()
+        fig = Figure(figsize=(5, 4), dpi=100)
+        self.axis = fig.add_subplot(111)
+        canvas = FigureCanvas(fig)  # a Gtk.DrawingArea
+        canvas.set_vexpand(True)
+        canvas.set_size_request(800, 250)
+        self.add(canvas)
 
-        cr.set_source_rgb(0, 0, 0)
-        cr.rectangle(-1.0 * SCREEN_SIZE, -0.5 * SCREEN_SIZE, SCREEN_SIZE,
-                     SCREEN_SIZE * 0.6)
-        #Axis
-        cr.move_to(-1.0 * SCREEN_SIZE + AXIS_MARGIN_SPACE,
-                   -0.5 * SCREEN_SIZE + AXIS_MARGIN_SPACE)  # Y
-        cr.line_to(-1.0 * SCREEN_SIZE + AXIS_MARGIN_SPACE,
-                   0.1 * SCREEN_SIZE - 10)
+    def recalculate_graph(self, points):
+        if not points.getLibsplines(): return
 
-        cr.move_to(-1.0 * SCREEN_SIZE + AXIS_MARGIN_SPACE,
-                   -0.5 * SCREEN_SIZE + AXIS_MARGIN_SPACE)  # X
-        cr.line_to(-10, -0.5 * SCREEN_SIZE + AXIS_MARGIN_SPACE)
-        cr.stroke()
+        # set the size of a timestep
+        dt = 0.00505
 
-        skip = 2
-        dT = 0.00505
-        start = AXIS_MARGIN_SPACE - SCREEN_SIZE
-        end = -2.0 * AXIS_MARGIN_SPACE
-        height = 0.5 * (SCREEN_SIZE) - AXIS_MARGIN_SPACE
-        zero = AXIS_MARGIN_SPACE - SCREEN_SIZE / 2.0
-        legend_entries = {}
-        if mypoints.getLibsplines():
-            distanceSpline = DistanceSpline(mypoints.getLibsplines())
-            traj = Trajectory(distanceSpline)
-            mypoints.addConstraintsToTrajectory(traj)
-            traj.Plan()
-            XVA = traj.GetPlanXVA(dT)
-            if XVA is not None:
-                self.draw_x_axis(cr, start, height, zero, XVA, end)
-                self.drawVelocity(cr, XVA, start, height, skip, zero, end,
-                                  legend_entries)
-                self.drawAcceleration(cr, XVA, start, height, skip, zero,
-                                      AXIS_MARGIN_SPACE, end, legend_entries)
-                self.drawVoltage(cr, XVA, start, height, skip, traj, zero, end,
-                                 legend_entries)
-                cr.set_source_rgb(0, 0, 0)
-                cr.move_to(-1.0 * AXIS_MARGIN_SPACE, zero + height / 2.0)
-                cr.line_to(AXIS_MARGIN_SPACE - SCREEN_SIZE,
-                           zero + height / 2.0)
-                self.drawLegend(cr, XVA, start, height, skip, zero, end,
-                                legend_entries)
-        cr.stroke()
+        # call C++ wrappers to calculate the trajectory
+        distanceSpline = DistanceSpline(points.getLibsplines())
+        traj = Trajectory(distanceSpline)
+        points.addConstraintsToTrajectory(traj)
+        traj.Plan()
+        XVA = traj.GetPlanXVA(dt)
 
-    def connectLines(self, cr, points, color):
-        for i in range(0, len(points) - 1):
-            set_color(cr, color)
-            cr.move_to(points[i][0], points[i][1])
-            cr.line_to(points[i + 1][0], points[i + 1][1])
-            cr.stroke()
+        # extract values to be graphed
+        total_steps_taken = XVA.shape[1]
+        total_time = dt * total_steps_taken
+        time = np.arange(total_time, step=dt)
+        position, velocity, acceleration = XVA
+        left_voltage, right_voltage = zip(*(traj.Voltage(x) for x in position))
 
-    def draw_x_axis(self, cr, start, height, zero, xva, end):
-        total_time = 0.00505 * len(xva[0])
-        for k in np.linspace(0, 1, 11):
-            self.tickMark(cr,
-                          k * np.abs(start - end) + start, zero + height / 2.0,
-                          10, palette["BLACK"])
-            cr.move_to(k * np.abs(start - end) + start,
-                       10 + zero + height / 2.0)
-            txt_scale = SCREEN_SIZE / 1000.0
-            display_text(cr, str(round(k * total_time, 3)), txt_scale,
-                         txt_scale, 1.0 / txt_scale, 1.0 / txt_scale)
-            cr.stroke()
+        # update graph
+        self.axis.clear()
+        self.axis.plot(time, velocity)
+        self.axis.plot(time, acceleration)
+        self.axis.plot(time, left_voltage)
+        self.axis.plot(time, right_voltage)
+        self.axis.legend(
+            ["Velocity", "Acceleration", "Left Voltage", "Right Voltage"])
+        self.axis.xaxis.set_label_text("Time (sec)")
 
-    def tickMark(self, cr, x, y, height, COLOR):
-        # X, Y is in the middle of the tick mark
-        set_color(cr, COLOR)
-        cr.move_to(x, y + (height / 2))
-        cr.line_to(x, y - (height / 2))
-        cr.stroke()
-
-    def HtickMark(self, cr, x, y, width, COLOR):
-        # X, Y is in the middle of the tick mark
-        set_color(cr, COLOR)
-        cr.move_to(x + (width / 2), y)
-        cr.line_to(x - (width / 2), y)
-        cr.stroke()
-
-    def drawLegend(self, cr, xva, start, height, skip, zero, end,
-                   legend_entries):
-        step_size = (end - start) / len(legend_entries)
-        margin_under_x_axis = height * 0.1
-        for index, (name, color) in enumerate(legend_entries.items()):
-            set_color(cr, color)
-            cr.move_to(start + index * step_size, zero - margin_under_x_axis)
-            txt_scale = SCREEN_SIZE / 900.0
-            display_text(cr, name, txt_scale, txt_scale, 1.0 / txt_scale,
-                         1.0 / txt_scale)
-
-    def drawVelocity(self, cr, xva, start, height, skip, zero, end,
-                     legend_entries):
-        COLOR = palette["RED"]
-        velocity = xva[1]
-        n_timesteps = len(velocity)
-        max_v = np.amax(velocity)
-        spacing = np.abs(start - end) / float(n_timesteps)
-        scaler = height / max_v
-        cr.set_source_rgb(1, 0, 0)
-        points = []
-        for i in range(0, len(velocity)):
-            if i % skip == 0:
-                points.append([
-                    start + (i * spacing),
-                    zero + height / 2.0 + (velocity[i] * scaler / 2.0)
-                ])
-        self.connectLines(cr, points, COLOR)
-
-        # draw axes marking
-        for i in np.linspace(-1, 1, 11):
-            self.HtickMark(cr, start, zero + i * height / 2.0 + height / 2.0,
-                           10, palette["BLACK"])
-            cr.set_source_rgb(1, 0, 0)
-            cr.move_to(start + 5, zero + i * height / 2.0 + height / 2.0)
-            txt_scale = SCREEN_SIZE / 1000.0
-            display_text(cr, str(round(i * max_v, 2)), txt_scale, txt_scale,
-                         1.0 / txt_scale, 1.0 / txt_scale)
-            cr.stroke()
-
-        # add entry to legend
-        legend_entries["Velocity"] = COLOR
-
-    def drawAcceleration(self, cr, xva, start, height, skip, zero, margin, end,
-                         legend_entries):
-        COLOR = palette["BLUE"]
-        accel = xva[2]
-        max_a = np.amax(accel)
-        min_a = np.amin(accel)
-        n_timesteps = len(accel)
-        spacing = np.abs(start - end) / float(n_timesteps)
-        scaler = height / (max_a - min_a)
-        cr.set_source_rgb(1, 0, 0)
-        points = []
-        for i in range(0, len(accel)):
-            if i % skip == 0:
-                points.append([
-                    start + (i * spacing), zero + ((accel[i] - min_a) * scaler)
-                ])
-        self.connectLines(cr, points, COLOR)
-
-        # draw axes marking
-        for i in np.linspace(0, 1, 11):
-            self.HtickMark(cr, -1.5 * margin, zero + i * height, 10,
-                           palette["BLACK"])
-            cr.set_source_rgb(0, 0, 1)
-            cr.move_to(-1.2 * margin, zero + i * height)
-            txt_scale = SCREEN_SIZE / 1000.0
-            display_text(cr, str(round(i * (max_a - min_a) + min_a,
-                                       2)), txt_scale, txt_scale,
-                         1.0 / txt_scale, 1.0 / txt_scale)
-            cr.stroke()
-
-        # draw legend
-        legend_entries["Acceleration"] = COLOR
-
-    def drawVoltage(self, cr, xva, start, height, skip, traj, zero, end,
-                    legend_entries):
-        COLOR_LEFT = palette["GREEN"]
-        COLOR_RIGHT = palette["CYAN"]
-        poses = xva[0]
-        n_timesteps = len(poses)
-        spacing = np.abs(start - end) / float(n_timesteps)
-        points_left = []
-        points_right = []
-        for i in range(0, len(poses)):
-            if i % skip == 0:
-                # libspline says the order is left-right
-                voltage = traj.Voltage(poses[i])
-                points_left.append([
-                    start + (i * spacing),
-                    zero + height / 2 + height * (voltage[0] / 24.0)
-                ])
-                points_right.append([
-                    start + (i * spacing),
-                    zero + height / 2 + height * (voltage[1] / 24.0)
-                ])
-        self.connectLines(cr, points_left, COLOR_LEFT)
-        self.connectLines(cr, points_right, COLOR_RIGHT)
-
-        for i in np.linspace(-1, 1, 7):
-            self.HtickMark(cr, -1.0 * SCREEN_SIZE,
-                           zero + i * height / 2.0 + height / 2.0, 10,
-                           palette["BLACK"])
-            cr.set_source_rgb(0, 1, 1)
-            cr.move_to(-1.0 * SCREEN_SIZE,
-                       zero + i * height / 2.0 + height / 2.0)
-            txt_scale = SCREEN_SIZE / 1000.0
-            display_text(cr, str(round(i * 12.0, 2)), txt_scale, txt_scale,
-                         1.0 / txt_scale, 1.0 / txt_scale)
-            cr.stroke()
-
-        legend_entries["Left Voltage"] = COLOR_LEFT
-        legend_entries["Right Voltage"] = COLOR_RIGHT
+        # renumber the x-axis to include the last point,
+        # the total time to drive the spline
+        self.axis.xaxis.set_ticks(np.linspace(0, total_time, num=8))