blob: ea6f5b2cef250801b3dd4aa88b5f6673499a9abb [file] [log] [blame]
Maxwell Henderson7af00982023-02-04 12:42:07 -08001#!/usr/bin/python3
2
3from __future__ import print_function
Maxwell Henderson765567c2023-03-05 14:22:17 -08004# matplotlib overrides fontconfig locations, so it needs to be imported before gtk.
5import matplotlib.pyplot as plt
Maxwell Hendersonc7bdf9d2023-03-23 10:15:31 -07006from matplotlib.backends.backend_gtk3agg import (FigureCanvasGTK3Agg as
7 FigureCanvas)
Maxwell Henderson7af00982023-02-04 12:42:07 -08008import os
9from frc971.control_loops.python import basic_window
10from frc971.control_loops.python.color import Color, palette
11import random
12import gi
milind-u18a901d2023-02-17 21:51:55 -080013import numpy as np
Maxwell Henderson7af00982023-02-04 12:42:07 -080014
15gi.require_version('Gtk', '3.0')
16from gi.repository import Gdk, Gtk
17import cairo
Maxwell Henderson165b17e2023-03-09 17:14:04 -080018from y2023.control_loops.python.graph_tools import to_theta, to_xy, alpha_blend, shift_angles, get_xy
milind-u600738b2023-02-22 14:42:19 -080019from y2023.control_loops.python.graph_tools import l1, l2, joint_center
20from y2023.control_loops.python.graph_tools import DRIVER_CAM_POINTS
21from y2023.control_loops.python import graph_paths
Maxwell Henderson7af00982023-02-04 12:42:07 -080022
Maxwell Henderson83cf6d62023-02-10 20:29:26 -080023from frc971.control_loops.python.basic_window import quit_main_loop, set_color, OverrideMatrix, identity
Maxwell Henderson7af00982023-02-04 12:42:07 -080024
25import shapely
26from shapely.geometry import Polygon
27
Maxwell Hendersonbfaf5ef2023-03-22 13:31:32 -070028from frc971.control_loops.python.constants import *
29
Maxwell Henderson7af00982023-02-04 12:42:07 -080030
Maxwell Henderson83cf6d62023-02-10 20:29:26 -080031def px(cr):
32 return OverrideMatrix(cr, identity)
33
34
35# Draw lines to cr + stroke.
36def draw_lines(cr, lines):
37 cr.move_to(lines[0][0], lines[0][1])
38 for pt in lines[1:]:
39 cr.line_to(pt[0], pt[1])
40 with px(cr):
41 cr.stroke()
42
43
Maxwell Henderson7af00982023-02-04 12:42:07 -080044def draw_px_cross(cr, length_px):
45 """Draws a cross with fixed dimensions in pixel space."""
46 with px(cr):
47 x, y = cr.get_current_point()
48 cr.move_to(x, y - length_px)
49 cr.line_to(x, y + length_px)
50 cr.stroke()
51
52 cr.move_to(x - length_px, y)
53 cr.line_to(x + length_px, y)
54 cr.stroke()
55
56
57def angle_dist_sqr(a1, a2):
58 """Distance between two points in angle space."""
59 return (a1[0] - a2[0])**2 + (a1[1] - a2[1])**2
60
61
62# Find the highest y position that intersects the vertical line defined by x.
63def inter_y(x):
milind-u18a901d2023-02-17 21:51:55 -080064 return np.sqrt((l2 + l1)**2 - (x - joint_center[0])**2) + joint_center[1]
Maxwell Henderson7af00982023-02-04 12:42:07 -080065
66
67# Define min and max l1 angles based on vertical constraints.
68def get_angle(boundary):
milind-u18a901d2023-02-17 21:51:55 -080069 h = np.sqrt((l1)**2 - (boundary - joint_center[0])**2) + joint_center[1]
70 return np.arctan2(h, boundary - joint_center[0])
Maxwell Henderson7af00982023-02-04 12:42:07 -080071
72
Maxwell Henderson7af00982023-02-04 12:42:07 -080073# Rotate a rasterized loop such that it aligns to when the parameters loop
74def rotate_to_jump_point(points):
75 last_pt = points[0]
76 for pt_i in range(1, len(points)):
77 pt = points[pt_i]
78 delta = last_pt[1] - pt[1]
milind-u18a901d2023-02-17 21:51:55 -080079 if abs(delta) > np.pi:
Maxwell Henderson7af00982023-02-04 12:42:07 -080080 return points[pt_i:] + points[:pt_i]
81 last_pt = pt
82 return points
83
84
85# shift points vertically by dy.
86def y_shift(points, dy):
87 return [(x, y + dy) for x, y in points]
88
89
Maxwell Henderson7af00982023-02-04 12:42:07 -080090# Get the closest point to a line from a test pt.
91def get_closest(prev, cur, pt):
92 dx_ang = (cur[0] - prev[0])
93 dy_ang = (cur[1] - prev[1])
94
milind-u18a901d2023-02-17 21:51:55 -080095 d = np.sqrt(dx_ang**2 + dy_ang**2)
Maxwell Henderson7af00982023-02-04 12:42:07 -080096 if (d < 0.000001):
milind-u18a901d2023-02-17 21:51:55 -080097 return prev, np.sqrt((prev[0] - pt[0])**2 + (prev[1] - pt[1])**2)
Maxwell Henderson7af00982023-02-04 12:42:07 -080098
99 pdx = -dy_ang / d
100 pdy = dx_ang / d
101
102 dpx = pt[0] - prev[0]
103 dpy = pt[1] - prev[1]
104
105 alpha = (dx_ang * dpx + dy_ang * dpy) / d / d
106
107 if (alpha < 0):
milind-u18a901d2023-02-17 21:51:55 -0800108 return prev, np.sqrt((prev[0] - pt[0])**2 + (prev[1] - pt[1])**2)
Maxwell Henderson7af00982023-02-04 12:42:07 -0800109 elif (alpha > 1):
milind-u18a901d2023-02-17 21:51:55 -0800110 return cur, np.sqrt((cur[0] - pt[0])**2 + (cur[1] - pt[1])**2)
Maxwell Henderson7af00982023-02-04 12:42:07 -0800111 else:
112 return (alpha_blend(prev[0], cur[0], alpha), alpha_blend(prev[1], cur[1], alpha)), \
113 abs(dpx * pdx + dpy * pdy)
114
115
116def closest_segment(lines, pt):
117 c_pt, c_pt_dist = get_closest(lines[-1], lines[0], pt)
118 for i in range(1, len(lines)):
119 prev = lines[i - 1]
120 cur = lines[i]
121 c_pt_new, c_pt_new_dist = get_closest(prev, cur, pt)
122 if c_pt_new_dist < c_pt_dist:
123 c_pt = c_pt_new
124 c_pt_dist = c_pt_new_dist
125 return c_pt, c_pt_dist
126
127
Maxwell Henderson4b8d02a2023-02-24 15:53:26 -0800128# Defining outline of robot
129class RobotOutline():
130
131 def __init__(self, drivetrain_pos, drivetrain_width, joint_center_radius,
132 joint_tower_pos, joint_tower_width, joint_tower_height,
133 driver_cam_pos, driver_cam_width, driver_cam_height):
134 self.drivetrain_pos = drivetrain_pos
135 self.drivetrain_width = drivetrain_width
136
137 self.joint_center_radius = joint_center_radius
138 self.joint_tower_pos = joint_tower_pos
139 self.joint_tower_width = joint_tower_width
140 self.joint_tower_height = joint_tower_height
141
142 self.driver_cam_pos = driver_cam_pos
143 self.driver_cam_width = driver_cam_width
144 self.driver_cam_height = driver_cam_height
145
146 def draw(self, cr):
147 set_color(cr, palette["BLUE"])
148
149 cr.move_to(self.drivetrain_pos[0], self.drivetrain_pos[1])
150 cr.line_to(self.drivetrain_pos[0] + self.drivetrain_width,
151 self.drivetrain_pos[1])
152
153 with px(cr):
154 cr.stroke()
155
156 # Draw joint center
157 cr.arc(joint_center[0], joint_center[1], self.joint_center_radius, 0,
158 2.0 * np.pi)
159 with px(cr):
160 cr.stroke()
161
162 # Draw joint tower
163 cr.rectangle(self.joint_tower_pos[0], self.joint_tower_pos[1],
164 self.joint_tower_width, self.joint_tower_height)
165
166 with px(cr):
167 cr.stroke()
168
169 # Draw driver cam
170 cr.set_source_rgba(1, 0, 0, 0.5)
171 cr.rectangle(self.driver_cam_pos[0], self.driver_cam_pos[1],
172 self.driver_cam_width, self.driver_cam_height)
173
174 with px(cr):
175 cr.fill()
176
177 def draw_theta(self, cr):
178 # TOOD(Max): add theta mode drawing
179 pass
180
181
182DRIVETRAIN_X = -0.490
183DRIVETRAIN_Y = 0.184
184DRIVETRAIN_WIDTH = 0.980
185
186JOINT_CENTER_RADIUS = 0.173 / 2
187
188JOINT_TOWER_X = -0.252
189JOINT_TOWER_Y = DRIVETRAIN_Y
190JOINT_TOWER_WIDTH = 0.098
191JOINT_TOWER_HEIGHT = 0.864
192
193DRIVER_CAM_X = DRIVER_CAM_POINTS[0][0]
194DRIVER_CAM_Y = DRIVER_CAM_POINTS[0][1]
195DRIVER_CAM_WIDTH = DRIVER_CAM_POINTS[-1][0] - DRIVER_CAM_POINTS[0][0]
196DRIVER_CAM_HEIGHT = DRIVER_CAM_POINTS[-1][1] - DRIVER_CAM_POINTS[0][1]
197
198
Maxwell Henderson0b5f1d12023-03-08 20:01:35 -0800199class SegmentSelector(basic_window.BaseWindow):
200
201 def __init__(self, segments):
202 super(SegmentSelector, self).__init__()
203
204 self.window = Gtk.Window()
205 self.window.set_title("Segment Selector")
206
207 self.segments = segments
208
209 self.segment_store = Gtk.ListStore(int, str)
210
211 for i, segment in enumerate(segments):
212 self.segment_store.append([i, segment.name])
213
214 self.segment_box = Gtk.ComboBox.new_with_model_and_entry(
215 self.segment_store)
216 self.segment_box.connect("changed", self.on_combo_changed)
217 self.segment_box.set_entry_text_column(1)
218
219 self.current_path_index = None
220
221 self.window.add(self.segment_box)
222 self.window.show_all()
223
224 def on_combo_changed(self, combo):
225 iter = combo.get_active_iter()
226
227 if iter is not None:
228 model = combo.get_model()
229 id, name = model[iter][:2]
230 print("Selected: ID=%d, name=%s" % (id, name))
231 self.current_path_index = id
232
233
Maxwell Hendersonc7bdf9d2023-03-23 10:15:31 -0700234ARM_AREA_WIDTH = 2 * (SCREEN_SIZE - 200)
235ARM_AREA_HEIGHT = SCREEN_SIZE
236
237
Maxwell Henderson7af00982023-02-04 12:42:07 -0800238# Create a GTK+ widget on which we will draw using Cairo
Maxwell Hendersonbfaf5ef2023-03-22 13:31:32 -0700239class ArmUi(Gtk.DrawingArea):
Maxwell Henderson7af00982023-02-04 12:42:07 -0800240
Maxwell Henderson0b5f1d12023-03-08 20:01:35 -0800241 def __init__(self, segments):
milind-u18a901d2023-02-17 21:51:55 -0800242 super(ArmUi, self).__init__()
Maxwell Henderson7af00982023-02-04 12:42:07 -0800243
Maxwell Hendersonc7bdf9d2023-03-23 10:15:31 -0700244 self.set_size_request(ARM_AREA_WIDTH, ARM_AREA_HEIGHT)
Maxwell Hendersonbfaf5ef2023-03-22 13:31:32 -0700245 self.center = (0, 0)
Maxwell Hendersonc7bdf9d2023-03-23 10:15:31 -0700246 self.shape = (ARM_AREA_WIDTH, ARM_AREA_HEIGHT)
Maxwell Henderson7af00982023-02-04 12:42:07 -0800247 self.theta_version = False
Maxwell Hendersonbfaf5ef2023-03-22 13:31:32 -0700248
249 self.init_extents()
250
251 self.connect('draw', self.on_draw)
Maxwell Henderson7af00982023-02-04 12:42:07 -0800252
Austin Schuh9a11ebd2023-02-26 14:16:31 -0800253 self.last_pos = to_xy(*graph_paths.points['Neutral'][:2])
milind-u600738b2023-02-22 14:42:19 -0800254 self.circular_index_select = 1
Maxwell Henderson7af00982023-02-04 12:42:07 -0800255
256 # Extra stuff for drawing lines.
Maxwell Henderson0b5f1d12023-03-08 20:01:35 -0800257 self.segments = segments
Maxwell Henderson7af00982023-02-04 12:42:07 -0800258 self.prev_segment_pt = None
259 self.now_segment_pt = None
260 self.spline_edit = 0
261 self.edit_control1 = True
262
milind-uadd8fa32023-02-24 23:37:36 -0800263 self.joint_thetas = None
264 self.joint_points = None
milind-u18a901d2023-02-17 21:51:55 -0800265 self.fig = plt.figure()
milind-uadd8fa32023-02-24 23:37:36 -0800266 self.axes = [
267 self.fig.add_subplot(3, 1, 1),
268 self.fig.add_subplot(3, 1, 2),
269 self.fig.add_subplot(3, 1, 3)
270 ]
271 self.fig.subplots_adjust(hspace=1.0)
milind-u18a901d2023-02-17 21:51:55 -0800272
Maxwell Hendersonf5c08a52023-02-23 17:28:35 -0800273 self.index = 0
274
Maxwell Henderson4b8d02a2023-02-24 15:53:26 -0800275 self.outline = RobotOutline([DRIVETRAIN_X, DRIVETRAIN_Y],
276 DRIVETRAIN_WIDTH, JOINT_CENTER_RADIUS,
277 [JOINT_TOWER_X, JOINT_TOWER_Y],
278 JOINT_TOWER_WIDTH, JOINT_TOWER_HEIGHT,
279 [DRIVER_CAM_X, DRIVER_CAM_Y],
280 DRIVER_CAM_WIDTH, DRIVER_CAM_HEIGHT)
281
Maxwell Hendersonc088b842023-03-11 07:22:02 -0800282 self.show_indicators = True
Maxwell Henderson4e57e0c2023-03-11 12:25:04 -0800283 # Lets you only view selected path
284 self.view_current = False
Maxwell Hendersonc088b842023-03-11 07:22:02 -0800285
Maxwell Henderson4890c202023-03-21 12:09:04 -0700286 self.editing = True
287
Maxwell Hendersonbfaf5ef2023-03-22 13:31:32 -0700288 self.x_offset = 0
289 self.y_offset = 0
290
Maxwell Henderson7af00982023-02-04 12:42:07 -0800291 def _do_button_press_internal(self, event):
292 o_x = event.x
293 o_y = event.y
Maxwell Hendersonbfaf5ef2023-03-22 13:31:32 -0700294 event.y -= self.y_offset
295 event.x -= self.x_offset
Maxwell Henderson7af00982023-02-04 12:42:07 -0800296 x = event.x - self.window_shape[0] / 2
297 y = self.window_shape[1] / 2 - event.y
298 scale = self.get_current_scale()
299 event.x = x / scale + self.center[0]
300 event.y = y / scale + self.center[1]
301 self.do_button_press(event)
302 event.x = o_x
303 event.y = o_y
304
Maxwell Hendersonbfaf5ef2023-03-22 13:31:32 -0700305 def init_extents(self):
Maxwell Henderson7af00982023-02-04 12:42:07 -0800306 if self.theta_version:
milind-u18a901d2023-02-17 21:51:55 -0800307 self.extents_x_min = -np.pi * 2
308 self.extents_x_max = np.pi * 2
309 self.extents_y_min = -np.pi * 2
310 self.extents_y_max = np.pi * 2
Maxwell Henderson7af00982023-02-04 12:42:07 -0800311 else:
312 self.extents_x_min = -40.0 * 0.0254
313 self.extents_x_max = 40.0 * 0.0254
314 self.extents_y_min = -4.0 * 0.0254
315 self.extents_y_max = 110.0 * 0.0254
316
Maxwell Hendersonbfaf5ef2023-03-22 13:31:32 -0700317 self.center = (0.5 * (self.extents_x_min + self.extents_x_max),
318 0.5 * (self.extents_y_max + self.extents_y_min))
319 self.shape = (1.0 * (self.extents_x_max - self.extents_x_min),
320 1.0 * (self.extents_y_max - self.extents_y_min))
321
322 def get_current_scale(self):
323 w_w, w_h = self.window_shape
324 w, h = self.shape
325 return min((w_w / w), (w_h / h))
326
327 def on_draw(self, widget, event):
328 cr = self.get_window().cairo_create()
329
330 self.window_shape = (self.get_window().get_geometry().width,
331 self.get_window().get_geometry().height)
332
333 cr.save()
334 cr.set_font_size(20)
335 cr.translate(self.window_shape[0] / 2, self.window_shape[1] / 2)
336 scale = self.get_current_scale()
337 cr.scale(scale, -scale)
338 cr.translate(-self.center[0], -self.center[1])
339 cr.reset_clip()
340 self.handle_draw(cr)
341 cr.restore()
342
343 def method_connect(self, event, cb):
344
345 def handler(obj, *args):
346 cb(*args)
347
348 self.window.connect(event, handler)
Maxwell Henderson7af00982023-02-04 12:42:07 -0800349
350 # Handle the expose-event by drawing
351 def handle_draw(self, cr):
Maxwell Henderson7af00982023-02-04 12:42:07 -0800352 # Fill the background color of the window with grey
353 set_color(cr, palette["GREY"])
354 cr.paint()
355
356 # Draw a extents rectangle
357 set_color(cr, palette["WHITE"])
358 cr.rectangle(self.extents_x_min, self.extents_y_min,
359 (self.extents_x_max - self.extents_x_min),
360 self.extents_y_max - self.extents_y_min)
361 cr.fill()
362
Maxwell Henderson93380322023-02-04 16:31:54 -0800363 if self.theta_version:
364 # Draw a filled white rectangle.
365 set_color(cr, palette["WHITE"])
milind-u18a901d2023-02-17 21:51:55 -0800366 cr.rectangle(-np.pi, -np.pi, np.pi * 2.0, np.pi * 2.0)
Maxwell Henderson93380322023-02-04 16:31:54 -0800367 cr.fill()
368
369 set_color(cr, palette["BLUE"])
370 for i in range(-6, 6):
milind-u18a901d2023-02-17 21:51:55 -0800371 cr.move_to(-40, -40 + i * np.pi)
372 cr.line_to(40, 40 + i * np.pi)
Maxwell Henderson93380322023-02-04 16:31:54 -0800373 with px(cr):
374 cr.stroke()
375
Maxwell Henderson93380322023-02-04 16:31:54 -0800376 set_color(cr, Color(0.0, 1.0, 0.2))
377 cr.move_to(self.last_pos[0], self.last_pos[1])
378 draw_px_cross(cr, 5)
379
Maxwell Henderson93380322023-02-04 16:31:54 -0800380 else:
Maxwell Henderson7af00982023-02-04 12:42:07 -0800381 # Draw a filled white rectangle.
382 set_color(cr, palette["WHITE"])
383 cr.rectangle(-2.0, -2.0, 4.0, 4.0)
384 cr.fill()
385
Maxwell Henderson4b8d02a2023-02-24 15:53:26 -0800386 self.outline.draw(cr)
milind-u18a901d2023-02-17 21:51:55 -0800387
388 # Draw max radius
Maxwell Henderson7af00982023-02-04 12:42:07 -0800389 set_color(cr, palette["BLUE"])
milind-u18a901d2023-02-17 21:51:55 -0800390 cr.arc(joint_center[0], joint_center[1], l2 + l1, 0, 2.0 * np.pi)
391 with px(cr):
392 cr.stroke()
393 cr.arc(joint_center[0], joint_center[1], l1 - l2, 0, 2.0 * np.pi)
Maxwell Henderson7af00982023-02-04 12:42:07 -0800394 with px(cr):
395 cr.stroke()
396
milind-u18a901d2023-02-17 21:51:55 -0800397 set_color(cr, Color(0.5, 1.0, 1))
Maxwell Henderson7af00982023-02-04 12:42:07 -0800398
milind-u600738b2023-02-22 14:42:19 -0800399 if not self.theta_version:
400 theta1, theta2 = to_theta(self.last_pos,
401 self.circular_index_select)
402 x, y = joint_center[0], joint_center[1]
403 cr.move_to(x, y)
404
405 x += np.cos(theta1) * l1
406 y += np.sin(theta1) * l1
407 cr.line_to(x, y)
408 x += np.cos(theta2) * l2
409 y += np.sin(theta2) * l2
410 cr.line_to(x, y)
411 with px(cr):
412 cr.stroke()
413
414 cr.move_to(self.last_pos[0], self.last_pos[1])
415 set_color(cr, Color(0.0, 1.0, 0.2))
416 draw_px_cross(cr, 20)
417
418 if self.theta_version:
419 set_color(cr, Color(0.0, 1.0, 0.2))
420 cr.move_to(self.last_pos[0], self.last_pos[1])
421 draw_px_cross(cr, 5)
422
Maxwell Henderson4b8d02a2023-02-24 15:53:26 -0800423 self.outline.draw_theta(cr)
milind-u600738b2023-02-22 14:42:19 -0800424
Maxwell Henderson7af00982023-02-04 12:42:07 -0800425 set_color(cr, Color(0.0, 0.5, 1.0))
Maxwell Henderson4e57e0c2023-03-11 12:25:04 -0800426 if not self.view_current:
427 for i in range(len(self.segments)):
428 color = None
429 if i == self.index:
430 continue
431 color = [0, random.random(), 1]
432 random.shuffle(color)
433 set_color(cr, Color(color[0], color[1], color[2]))
434 self.segments[i].DrawTo(cr, self.theta_version)
Maxwell Henderson165b17e2023-03-09 17:14:04 -0800435
Maxwell Henderson4e57e0c2023-03-11 12:25:04 -0800436 with px(cr):
437 cr.stroke()
Maxwell Henderson7af00982023-02-04 12:42:07 -0800438
Austin Schuhdba32702023-03-04 22:51:05 -0800439 # Draw current spline in black
440 color = [0, 0, 0]
441 set_color(cr, Color(color[0], color[1], color[2]))
442 self.segments[self.index].DrawTo(cr, self.theta_version)
Maxwell Henderson165b17e2023-03-09 17:14:04 -0800443
Maxwell Hendersonc088b842023-03-11 07:22:02 -0800444 with px(cr):
445 cr.stroke()
Maxwell Henderson165b17e2023-03-09 17:14:04 -0800446 control1 = get_xy(self.segments[self.index].control1)
447 control2 = get_xy(self.segments[self.index].control2)
448
449 if self.theta_version:
450 control1 = shift_angles(self.segments[self.index].control1)
451 control2 = shift_angles(self.segments[self.index].control2)
452
Maxwell Hendersonc088b842023-03-11 07:22:02 -0800453 if self.show_indicators:
454 set_color(cr, Color(1.0, 0.0, 1.0))
455 cr.move_to(control1[0] + 0.02, control1[1])
456 cr.arc(control1[0], control1[1], 0.02, 0, 2.0 * np.pi)
457 with px(cr):
458 cr.stroke()
459 set_color(cr, Color(1.0, 0.7, 0.0))
460 cr.move_to(control2[0] + 0.02, control2[1])
461 cr.arc(control2[0], control2[1], 0.02, 0, 2.0 * np.pi)
Maxwell Henderson165b17e2023-03-09 17:14:04 -0800462
Austin Schuhdba32702023-03-04 22:51:05 -0800463 with px(cr):
464 cr.stroke()
465
Maxwell Henderson7af00982023-02-04 12:42:07 -0800466 set_color(cr, Color(0.0, 1.0, 0.5))
milind-u18a901d2023-02-17 21:51:55 -0800467
milind-uadd8fa32023-02-24 23:37:36 -0800468 # Create the plots
469 if self.joint_thetas:
470 if self.joint_points:
471 titles = ["Proximal", "Distal", "Roll joint"]
472 for i in range(len(self.joint_points)):
473 self.axes[i].clear()
474 self.axes[i].plot(self.joint_thetas[0],
475 self.joint_thetas[1][i])
476 self.axes[i].scatter([self.joint_points[i][0]],
477 [self.joint_points[i][1]],
478 s=10,
479 c="red")
480 self.axes[i].set_title(titles[i])
481 plt.title("Joint Angle")
milind-u18a901d2023-02-17 21:51:55 -0800482 plt.xlabel("t (0 to 1)")
483 plt.ylabel("theta (rad)")
484
485 self.fig.canvas.draw()
Maxwell Henderson7af00982023-02-04 12:42:07 -0800486
487 def cur_pt_in_theta(self):
milind-u18a901d2023-02-17 21:51:55 -0800488 if self.theta_version: return self.last_pos
milind-u600738b2023-02-22 14:42:19 -0800489 return to_theta(self.last_pos,
490 self.circular_index_select,
491 cross_point=-np.pi,
492 die=False)
Maxwell Henderson7af00982023-02-04 12:42:07 -0800493
milind-u18a901d2023-02-17 21:51:55 -0800494 def do_motion(self, event):
495 o_x = event.x
496 o_y = event.y
Maxwell Hendersonbfaf5ef2023-03-22 13:31:32 -0700497 event.x -= self.x_offset
498 event.y -= self.y_offset
milind-u18a901d2023-02-17 21:51:55 -0800499 x = event.x - self.window_shape[0] / 2
500 y = self.window_shape[1] / 2 - event.y
501 scale = self.get_current_scale()
502 event.x = x / scale + self.center[0]
503 event.y = y / scale + self.center[1]
504
milind-u52af04c2023-02-25 14:03:08 -0800505 segment = self.segments[self.index]
506 self.joint_thetas = segment.joint_thetas()
milind-u18a901d2023-02-17 21:51:55 -0800507
milind-u52af04c2023-02-25 14:03:08 -0800508 hovered_t = segment.intersection(event)
509 if hovered_t:
510 min_diff = np.inf
511 closest_t = None
512 closest_thetas = None
513 for i in range(len(self.joint_thetas[0])):
514 t = self.joint_thetas[0][i]
515 diff = abs(t - hovered_t)
516 if diff < min_diff:
517 min_diff = diff
518 closest_t = t
519 closest_thetas = [
520 self.joint_thetas[1][0][i], self.joint_thetas[1][1][i],
521 self.joint_thetas[1][2][i]
522 ]
523 self.joint_points = [(closest_t, closest_theta)
524 for closest_theta in closest_thetas]
milind-u18a901d2023-02-17 21:51:55 -0800525
526 event.x = o_x
527 event.y = o_y
528
Maxwell Hendersonbfaf5ef2023-03-22 13:31:32 -0700529 self.queue_draw()
Maxwell Henderson7af00982023-02-04 12:42:07 -0800530
Maxwell Henderson75c70122023-03-22 20:16:47 -0700531 def switch_theta(self):
532 # Toggle between theta and xy renderings
533 if self.theta_version:
534 theta1, theta2 = self.last_pos
535 data = to_xy(theta1, theta2)
536 self.circular_index_select = int(
537 np.floor((theta2 - theta1) / np.pi))
538 self.last_pos = (data[0], data[1])
539 else:
540 self.last_pos = self.cur_pt_in_theta()
541
542 self.theta_version = not self.theta_version
543 self.init_extents()
544
Maxwell Henderson7af00982023-02-04 12:42:07 -0800545 def do_key_press(self, event):
546 keyval = Gdk.keyval_to_lower(event.keyval)
547 print("Gdk.KEY_" + Gdk.keyval_name(keyval))
548 if keyval == Gdk.KEY_q:
549 print("Found q key and exiting.")
550 quit_main_loop()
551 elif keyval == Gdk.KEY_c:
552 # Increment which arm solution we render
553 self.circular_index_select += 1
554 print(self.circular_index_select)
555 elif keyval == Gdk.KEY_v:
556 # Decrement which arm solution we render
557 self.circular_index_select -= 1
558 print(self.circular_index_select)
Maxwell Henderson7af00982023-02-04 12:42:07 -0800559
560 elif keyval == Gdk.KEY_r:
561 self.prev_segment_pt = self.now_segment_pt
562
Maxwell Hendersonf5c08a52023-02-23 17:28:35 -0800563 elif keyval == Gdk.KEY_o:
564 # Only prints current segment
565 print(repr(self.segments[self.index]))
Maxwell Henderson7af00982023-02-04 12:42:07 -0800566 elif keyval == Gdk.KEY_g:
567 # Generate theta points.
568 if self.segments:
Maxwell Hendersonf5c08a52023-02-23 17:28:35 -0800569 print(repr(self.segments[self.index].ToThetaPoints()))
Maxwell Henderson7af00982023-02-04 12:42:07 -0800570 elif keyval == Gdk.KEY_e:
571 best_pt = self.now_segment_pt
572 best_dist = 1e10
573 for segment in self.segments:
574 d = angle_dist_sqr(segment.start, self.now_segment_pt)
575 if (d < best_dist):
576 best_pt = segment.start
577 best_dist = d
578 d = angle_dist_sqr(segment.end, self.now_segment_pt)
579 if (d < best_dist):
580 best_pt = segment.end
581 best_dist = d
582 self.now_segment_pt = best_pt
583
Austin Schuha09e1bb2023-02-25 22:34:41 -0800584 elif keyval == Gdk.KEY_p:
585 if self.index > 0:
586 self.index -= 1
587 else:
588 self.index = len(self.segments) - 1
589 print("Switched to segment:", self.segments[self.index].name)
Austin Schuh9a11ebd2023-02-26 14:16:31 -0800590 self.segments[self.index].Print(graph_paths.points)
Austin Schuha09e1bb2023-02-25 22:34:41 -0800591
Maxwell Hendersonc088b842023-03-11 07:22:02 -0800592 elif keyval == Gdk.KEY_i:
593 self.show_indicators = not self.show_indicators
594
Austin Schuha09e1bb2023-02-25 22:34:41 -0800595 elif keyval == Gdk.KEY_n:
Maxwell Hendersonf5c08a52023-02-23 17:28:35 -0800596 self.index += 1
597 self.index = self.index % len(self.segments)
milind-u2a28c592023-02-24 23:13:04 -0800598 print("Switched to segment:", self.segments[self.index].name)
Austin Schuh9a11ebd2023-02-26 14:16:31 -0800599 self.segments[self.index].Print(graph_paths.points)
Maxwell Hendersonf5c08a52023-02-23 17:28:35 -0800600
Maxwell Henderson4890c202023-03-21 12:09:04 -0700601 elif keyval == Gdk.KEY_d:
602 self.editing = not self.editing
603
Maxwell Henderson4e57e0c2023-03-11 12:25:04 -0800604 elif keyval == Gdk.KEY_l:
605 self.view_current = not self.view_current
606
Maxwell Henderson7af00982023-02-04 12:42:07 -0800607 elif keyval == Gdk.KEY_t:
Maxwell Henderson75c70122023-03-22 20:16:47 -0700608 self.switch_theta()
Maxwell Henderson7af00982023-02-04 12:42:07 -0800609
610 elif keyval == Gdk.KEY_z:
611 self.edit_control1 = not self.edit_control1
612 if self.edit_control1:
Maxwell Hendersonf5c08a52023-02-23 17:28:35 -0800613 self.now_segment_pt = self.segments[self.index].control1
Maxwell Henderson7af00982023-02-04 12:42:07 -0800614 else:
Maxwell Hendersonf5c08a52023-02-23 17:28:35 -0800615 self.now_segment_pt = self.segments[self.index].control2
Maxwell Henderson7af00982023-02-04 12:42:07 -0800616 if not self.theta_version:
617 data = to_xy(self.now_segment_pt[0], self.now_segment_pt[1])
618 self.last_pos = (data[0], data[1])
619 else:
620 self.last_pos = self.now_segment_pt
621
622 print("self.last_pos: ", self.last_pos, " ci: ",
623 self.circular_index_select)
624
Maxwell Hendersonbfaf5ef2023-03-22 13:31:32 -0700625 self.queue_draw()
Maxwell Henderson7af00982023-02-04 12:42:07 -0800626
627 def do_button_press(self, event):
Maxwell Henderson4890c202023-03-21 12:09:04 -0700628
milind-u600738b2023-02-22 14:42:19 -0800629 last_pos = self.last_pos
Maxwell Henderson7af00982023-02-04 12:42:07 -0800630 self.last_pos = (event.x, event.y)
milind-u600738b2023-02-22 14:42:19 -0800631 pt_theta = self.cur_pt_in_theta()
632 if pt_theta is None:
633 self.last_pos = last_pos
634 return
635
636 self.now_segment_pt = np.array(shift_angles(pt_theta))
Maxwell Henderson7af00982023-02-04 12:42:07 -0800637
Maxwell Henderson4890c202023-03-21 12:09:04 -0700638 if self.editing:
639 if self.edit_control1:
640 self.segments[self.index].control1 = self.now_segment_pt
641 else:
642 self.segments[self.index].control2 = self.now_segment_pt
Maxwell Henderson7af00982023-02-04 12:42:07 -0800643
Austin Schuh9a11ebd2023-02-26 14:16:31 -0800644 print('Clicked at theta: np.array([%s, %s])' %
645 (self.now_segment_pt[0], self.now_segment_pt[1]))
Maxwell Henderson7af00982023-02-04 12:42:07 -0800646 if not self.theta_version:
Austin Schuh9a11ebd2023-02-26 14:16:31 -0800647 print(
648 'Clicked at to_theta_with_circular_index(%.3f, %.3f, circular_index=%d)'
649 % (self.last_pos[0], self.last_pos[1],
Maxwell Henderson7af00982023-02-04 12:42:07 -0800650 self.circular_index_select))
651
Austin Schuh9a11ebd2023-02-26 14:16:31 -0800652 self.segments[self.index].Print(graph_paths.points)
Maxwell Henderson7af00982023-02-04 12:42:07 -0800653
Maxwell Hendersonbfaf5ef2023-03-22 13:31:32 -0700654 self.queue_draw()
Maxwell Henderson7af00982023-02-04 12:42:07 -0800655
656
Maxwell Hendersonbfaf5ef2023-03-22 13:31:32 -0700657class Window(Gtk.Window):
658
659 def __init__(self, segments):
660 super().__init__(title="Drawing Area")
661
662 self.segment_store = Gtk.ListStore(int, str)
663
664 for i, segment in enumerate(segments):
665 self.segment_store.append([i, segment.name])
666
667 self.segment_box = Gtk.ComboBox.new_with_model_and_entry(
668 self.segment_store)
669 self.segment_box.connect("changed", self.on_combo_changed)
670 self.segment_box.set_entry_text_column(1)
671
672 self.arm_draw = ArmUi(segments)
673
674 self.arm_draw.y_offset = self.segment_box.get_allocation().width
675
676 print('Starting with segment: ',
677 self.arm_draw.segments[self.arm_draw.index].name)
678 self.arm_draw.segments[self.arm_draw.index].Print(graph_paths.points)
679
680 self.set_events(Gdk.EventMask.BUTTON_PRESS_MASK
681 | Gdk.EventMask.BUTTON_RELEASE_MASK
682 | Gdk.EventMask.POINTER_MOTION_MASK
683 | Gdk.EventMask.SCROLL_MASK
684 | Gdk.EventMask.KEY_PRESS_MASK)
685 self.method_connect('map-event', self.do_map_event)
686 self.method_connect("key-press-event", self.arm_draw.do_key_press)
687 self.method_connect("motion-notify-event", self.arm_draw.do_motion)
688 self.method_connect("button-press-event",
689 self.arm_draw._do_button_press_internal)
Maxwell Hendersonbfaf5ef2023-03-22 13:31:32 -0700690
691 self.grid = Gtk.Grid()
692 self.add(self.grid)
693
694 self.grid.attach(self.arm_draw, 0, 1, 1, 1)
695
Maxwell Henderson75c70122023-03-22 20:16:47 -0700696 self.isolate_button = Gtk.Button(label="Toggle Path Isolation")
697 self.isolate_button.connect('clicked', self.on_button_click)
698
699 self.theta_button = Gtk.Button(label="Toggle Theta Mode")
700 self.theta_button.connect('clicked', self.on_button_click)
701
702 self.editing_button = Gtk.Button(label="Toggle Editing Mode")
703 self.editing_button.connect('clicked', self.on_button_click)
704
705 self.indicator_button = Gtk.Button(
706 label="Toggle Control Point Indicators")
707 self.indicator_button.connect('clicked', self.on_button_click)
708
709 self.box = Gtk.Box(spacing=6)
710 self.grid.attach(self.box, 0, 0, 1, 1)
711
Maxwell Hendersonc7bdf9d2023-03-23 10:15:31 -0700712 self.figure_canvas = FigureCanvas(self.arm_draw.fig)
713 self.figure_canvas.set_size_request(500, 300)
714
715 self.grid.attach(self.figure_canvas, 1, 1, 1, 1)
716
Maxwell Henderson75c70122023-03-22 20:16:47 -0700717 self.box.pack_start(self.segment_box, False, False, 0)
718 self.box.pack_start(self.isolate_button, False, False, 0)
719 self.box.pack_start(self.theta_button, False, False, 0)
720 self.box.pack_start(self.editing_button, False, False, 0)
721 self.box.pack_start(self.indicator_button, False, False, 0)
Maxwell Hendersonbfaf5ef2023-03-22 13:31:32 -0700722
723 def on_combo_changed(self, combo):
724 iter = combo.get_active_iter()
725
726 if iter is not None:
727 model = combo.get_model()
728 id, name = model[iter][:2]
729 print("Selected: ID=%d, name=%s" % (id, name))
730 self.arm_draw.index = id
Maxwell Henderson75c70122023-03-22 20:16:47 -0700731 self.arm_draw.queue_draw()
Maxwell Hendersonbfaf5ef2023-03-22 13:31:32 -0700732
733 def method_connect(self, event, cb):
734
735 def handler(obj, *args):
736 cb(*args)
737
738 self.connect(event, handler)
739
740 def do_map_event(self, event):
Maxwell Henderson75c70122023-03-22 20:16:47 -0700741 self.arm_draw.y_offset = self.box.get_allocation().height
742
743 def on_button_click(self, button):
744 if self.isolate_button == button:
745 self.arm_draw.view_current = not self.arm_draw.view_current
746 elif self.theta_button == button:
747 self.arm_draw.switch_theta()
748 elif self.editing_button == button:
749 self.arm_draw.editing = not self.arm_draw.editing
750 elif self.indicator_button == button:
751 self.arm_draw.show_indicators = not self.arm_draw.show_indicators
752 self.arm_draw.queue_draw()
Maxwell Hendersonbfaf5ef2023-03-22 13:31:32 -0700753
754
755window = Window(graph_paths.segments)
756window.show_all()
Maxwell Henderson7af00982023-02-04 12:42:07 -0800757basic_window.RunApp()