frc971/control_loops/python/path_edit.py

#!/usr/bin/python3
from __future__ import print_function
import os
import sys
from color import palette
from graph import Graph
import gi
import numpy as np
gi.require_version('Gtk', '3.0')
from gi.repository import Gdk, Gtk, GLib
import cairo
from libspline import Spline, DistanceSpline
import enum
import json
import copy
from constants import FIELD
from constants import get_json_folder
from constants import ROBOT_SIDE_TO_BALL_CENTER, ROBOT_SIDE_TO_HATCH_PANEL, HATCH_PANEL_WIDTH, BALL_RADIUS
from drawing_constants import set_color, draw_px_cross, draw_px_x, display_text, draw_control_points
from multispline import Multispline, ControlPointIndex
import time


class Mode(enum.Enum):
    kViewing = 0
    kPlacing = 1
    kEditing = 2


class FieldWidget(Gtk.DrawingArea):
    """Create a GTK+ widget on which we will draw using Cairo"""
    def __init__(self):
        super(FieldWidget, self).__init__()
        self.set_field(FIELD)
        self.set_size_request(self.mToPx(self.field.width),
                              self.mToPx(self.field.length))

        self.multisplines = []
        self.graph = Graph()
        self.graph.cursor_watcher = self
        self.set_vexpand(True)
        self.set_hexpand(True)
        self.undo_history = []
        # init field drawing
        # add default spline for testing purposes
        # init editing / viewing modes and pointer location
        self.mode = Mode.kPlacing
        self.mousex = 0
        self.mousey = 0
        self.module_path = os.path.dirname(os.path.realpath(sys.argv[0]))
        self.path_to_export = os.path.join(self.module_path,
                                           'points_for_pathedit.json')

        # For the editing mode
        self.control_point_index = None
        self.active_multispline_index = -1

        self.zoom_transform = cairo.Matrix()

        self.set_events(Gdk.EventMask.BUTTON_PRESS_MASK
                        | Gdk.EventMask.BUTTON_PRESS_MASK
                        | Gdk.EventMask.BUTTON_RELEASE_MASK
                        | Gdk.EventMask.POINTER_MOTION_MASK
                        | Gdk.EventMask.SCROLL_MASK)

    @property
    def active_multispline(self):
        """Get the current active multispline or create a new one"""
        if not self.multisplines:
            self.multisplines.append(Multispline())
            self.active_multispline_index = -1

        return self.multisplines[self.active_multispline_index]

    def set_field(self, field):
        self.field = field
        try:
            self.field_png = cairo.ImageSurface.create_from_png(
                "frc971/control_loops/python/field_images/" +
                self.field.field_id + ".png")
        except cairo.Error:
            self.field_png = None

        self.queue_draw()

    def invert(self, transform):
        xx, yx, xy, yy, x0, y0 = transform
        matrix = cairo.Matrix(xx, yx, xy, yy, x0, y0)
        matrix.invert()
        return matrix

    # returns the transform from widget space to field space
    @property
    def input_transform(self):
        # the transform for input needs to be the opposite of the transform for drawing
        return self.invert(self.field_transform.multiply(self.zoom_transform))

    @property
    def field_transform(self):
        field_transform = cairo.Matrix()
        field_transform.scale(1, -1)  # flipped y-axis
        field_transform.scale(1 / self.pxToM_scale(), 1 / self.pxToM_scale())
        field_transform.translate(self.field.width / 2,
                                  -1 * self.field.length / 2)
        return field_transform

    # returns the scale from pixels in field space to meters in field space
    def pxToM_scale(self):
        available_space = self.get_allocation()
        return np.maximum(self.field.width / available_space.width,
                          self.field.length / available_space.height)

    def pxToM(self, p):
        return p * self.pxToM_scale()

    def mToPx(self, m):
        return m / self.pxToM_scale()

    def draw_robot_at_point(self, cr, spline, t):
        """Draws the robot at a point along a Spline or DistanceSpline"""

        # we accept both Spline and DistanceSpline
        if type(spline) is Spline:
            point = spline.Point(t)
            theta = spline.Theta(t)
        elif type(spline) is DistanceSpline:
            point = spline.XY(t)
            theta = spline.Theta(t)
        else:
            raise TypeError(
                f"expected Spline or DistanceSpline (got {type(spline)})")

        # Transform so that +y is forward along the spline
        transform = cairo.Matrix()
        transform.translate(*point)
        transform.rotate(theta - np.pi / 2)

        cr.save()
        cr.set_matrix(transform.multiply(cr.get_matrix()))

        # Draw Robot
        set_color(cr, palette["BLACK"])
        cr.rectangle(-self.field.robot.width / 2, -self.field.robot.length / 2,
                     self.field.robot.width, self.field.robot.length)
        cr.stroke()

        #Draw Ball
        set_color(cr, palette["ORANGE"], 0.5)
        cr.arc(0, self.field.robot.length / 2 + BALL_RADIUS, BALL_RADIUS, 0,
               2 * np.pi)
        cr.stroke()

        # undo the transform
        cr.restore()

    def do_draw(self, cr):  # main
        cr.set_matrix(
            self.field_transform.multiply(self.zoom_transform).multiply(
                cr.get_matrix()))

        cr.save()

        set_color(cr, palette["BLACK"])

        cr.set_line_width(self.pxToM(1))
        cr.rectangle(-0.5 * self.field.width, -0.5 * self.field.length,
                     self.field.width, self.field.length)
        cr.set_line_join(cairo.LINE_JOIN_ROUND)
        cr.stroke()

        if self.field_png:
            cr.save()
            cr.translate(-0.5 * self.field.width, 0.5 * self.field.length)
            cr.scale(
                self.field.width / self.field_png.get_width(),
                -self.field.length / self.field_png.get_height(),
            )
            cr.set_source_surface(self.field_png)
            cr.paint()
            cr.restore()

        # update everything

        cr.set_line_width(self.pxToM(1))
        if self.mode == Mode.kPlacing or self.mode == Mode.kViewing:
            set_color(cr, palette["BLACK"])
            for multispline in self.multisplines:
                for i, point in enumerate(multispline.staged_points):
                    draw_px_x(cr, point[0], point[1], self.pxToM(2))
            set_color(cr, palette["WHITE"])
        elif self.mode == Mode.kEditing:
            set_color(cr, palette["BLACK"])
            if len(self.multisplines) != 0 and self.multisplines[0].getSplines(
            ):
                self.draw_splines(cr)

            for multispline in self.multisplines:
                for i, points in enumerate(multispline.getSplines()):
                    points = [np.array([x, y]) for (x, y) in points]
                    draw_control_points(cr,
                                        points,
                                        width=self.pxToM(5),
                                        radius=self.pxToM(2))

                    p0, p1, p2, p3, p4, p5 = points
                    first_tangent = p0 + 2.0 * (p1 - p0)
                    second_tangent = p5 + 2.0 * (p4 - p5)
                    cr.set_source_rgb(0, 0.5, 0)
                    cr.move_to(*p0)
                    cr.set_line_width(self.pxToM(1.0))
                    cr.line_to(*first_tangent)
                    cr.move_to(*first_tangent)
                    cr.line_to(*p2)

                    cr.move_to(*p5)
                    cr.line_to(*second_tangent)

                    cr.move_to(*second_tangent)
                    cr.line_to(*p3)

                    cr.stroke()
                    cr.set_line_width(self.pxToM(2))
            set_color(cr, palette["WHITE"])

        cr.paint_with_alpha(0.2)

        draw_px_cross(cr, self.mousex, self.mousey, self.pxToM(2))
        cr.restore()

    def draw_splines(self, cr):
        for multispline in self.multisplines:
            for i, spline in enumerate(multispline.getLibsplines()):
                # draw lots of really small line segments to
                # approximate the shape of the spline
                for k in np.linspace(0.02, 1, 200):
                    cr.move_to(*spline.Point(k - 0.008))
                    cr.line_to(*spline.Point(k))
                    cr.stroke()

                if i == 0:
                    self.draw_robot_at_point(cr, spline, 0)
                self.draw_robot_at_point(cr, spline, 1)

        mouse = np.array((self.mousex, self.mousey))

        multispline, result = Multispline.nearest_distance(
            self.multisplines, mouse)

        # if the mouse is close enough, draw the robot
        if result is not None and result.fun < 2:
            distance_spline = DistanceSpline(multispline.getLibsplines())
            x = result.x[0]

            # draw the robot to show its width
            self.draw_robot_at_point(cr, distance_spline, x)

            multispline_index = self.multisplines.index(multispline)
            self.graph.place_cursor(multispline_index, distance=result.x[0])
        elif self.graph.cursor is not None:
            multispline_index, x = self.graph.find_cursor()
            distance_spline = DistanceSpline(
                self.multisplines[multispline_index].getLibsplines())

            self.draw_robot_at_point(cr, distance_spline, x)

            # clear the cursor each draw so it doesn't persist
            # after you move off the spline
            self.graph.cursor = None

    def export_json(self, file_name):
        self.path_to_export = os.path.join(
            self.module_path,  # position of the python
            "../../..",  # root of the repository
            get_json_folder(self.field),  # path from the root
            file_name  # selected file
        )

        # Will export to json file
        multisplines_object = [
            multispline.toJsonObject() for multispline in self.multisplines
        ]
        print(multisplines_object)
        with open(self.path_to_export, mode='w') as points_file:
            json.dump(multisplines_object, points_file)

    def import_json(self, file_name):
        self.path_to_export = os.path.join(
            self.module_path,  # position of the python
            "../../..",  # root of the repository
            get_json_folder(self.field),  # path from the root
            file_name  # selected file
        )

        # import from json file
        print("LOADING LOAD FROM " + file_name)  # Load takes a few seconds
        with open(self.path_to_export) as points_file:
            multisplines_object = json.load(points_file)

        self.attempt_append_multisplines()

        # TODO: Export multisplines in different files
        if type(multisplines_object) is dict:
            multisplines_object = [multisplines_object]
        else:
            self.multisplines = []

        # if people messed with the spline json,
        # it might not be the right length
        # so give them a nice error message
        for multispline_object in multisplines_object:
            print(multispline_object)
            try:  # try to salvage as many segments of the spline as possible
                self.multisplines.append(
                    Multispline.fromJsonObject(multispline_object))
            except IndexError:
                # check if they're both 6+5*(k-1) long
                expected_length = 6 + 5 * (multispline_object["spline_count"] -
                                           1)
                x_len = len(multispline_object["spline_x"])
                y_len = len(multispline_object["spline_x"])
                if x_len is not expected_length:
                    print(
                        "Error: spline x values were not the expected length; expected {} got {}"
                        .format(expected_length, x_len))
                elif y_len is not expected_length:
                    print(
                        "Error: spline y values were not the expected length; expected {} got {}"
                        .format(expected_length, y_len))

        print("SPLINES LOADED")
        self.mode = Mode.kEditing
        self.queue_draw()
        self.graph.schedule_recalculate(self.multisplines)

    def attempt_append_multisplines(self):
        if len(self.undo_history
               ) == 0 or self.multisplines != self.undo_history[-1]:
            self.undo_history.append(copy.deepcopy(self.multisplines))

    def clear(self, should_attempt_append=True):
        if should_attempt_append:
            self.attempt_append_multisplines()
        self.multisplines = []
        #recalulate graph using new points
        self.graph.axis.clear()
        self.graph.queue_draw()
        #allow placing again
        self.mode = Mode.kPlacing
        #redraw entire graph
        self.queue_draw()

    def undo(self):
        try:
            self.undo_history.pop()
        except IndexError:
            return
        if len(self.undo_history) == 0:
            self.clear(should_attempt_append=False)  #clear, don't do anything
            return
        if len(self.multisplines) > 0 and not any(
                multispline.staged_points
                for multispline in self.multisplines):
            self.mode = Mode.kEditing
        else:
            self.mode = Mode.kPlacing
            self.clear(should_attempt_append=False)
        self.multisplines = copy.deepcopy(self.undo_history[-1])
        self.queue_draw()

    def do_key_press_event(self, event):
        keyval = Gdk.keyval_to_lower(event.keyval)
        if keyval == Gdk.KEY_z and event.state & Gdk.ModifierType.CONTROL_MASK:
            self.undo()
        # TODO: This should be a button
        if keyval == Gdk.KEY_p:
            self.mode = Mode.kPlacing
            # F0 = A1
            # B1 = 2F0 - E0
            # C1= d0 + 4F0 - 4E0
            multispline = self.active_multispline
            multispline.extrapolate()
            self.queue_draw()
        elif keyval == Gdk.KEY_m:
            self.multisplines.append(Multispline())
            self.active_spline_index = len(self.multisplines) - 1
            self.mode = Mode.kPlacing

            multispline = self.multisplines[-2]
            #multispline.extrapolate()
            self.queue_draw()

    def do_button_release_event(self, event):
        self.attempt_append_multisplines()
        self.mousex, self.mousey = self.input_transform.transform_point(
            event.x, event.y)
        if self.mode == Mode.kEditing:
            if self.control_point_index != None:
                multispline = self.multisplines[
                    self.control_point_index.multispline_index]

                multispline.setControlPoint(self.control_point_index,
                                            self.mousex, self.mousey)

                multispline.splineExtrapolate(
                    self.control_point_index.spline_index)

                multispline.update_lib_spline()
                self.graph.schedule_recalculate(self.multisplines)

                self.control_point_index = None

    def do_button_press_event(self, event):
        self.mousex, self.mousey = self.input_transform.transform_point(
            event.x, event.y)

        if self.mode == Mode.kPlacing:
            if self.active_multispline.addPoint(self.mousex, self.mousey):
                self.mode = Mode.kEditing
                self.graph.schedule_recalculate(self.multisplines)
        elif self.mode == Mode.kEditing:
            # Now after we have no control point index,
            # the user can click for new point
            if self.control_point_index == None:
                # Get clicked point
                # Find nearest
                # Move nearest to clicked
                cur_p = [self.mousex, self.mousey]
                # Get the distance between each for x and y
                # Save the index of the point closest
                nearest = 1  # Max distance away a the selected point can be in meters
                index_of_closest = 0
                for index_multisplines, multispline in enumerate(
                        self.multisplines):
                    for index_splines, points in enumerate(
                            multispline.getSplines()):
                        for index_points, val in enumerate(points):
                            distance = np.sqrt((cur_p[0] - val[0])**2 +
                                               (cur_p[1] - val[1])**2)
                            if distance < nearest:
                                nearest = distance
                                index_of_closest = index_points
                                self.control_point_index = ControlPointIndex(
                                    index_multisplines, index_splines,
                                    index_points)
        self.queue_draw()

    def do_motion_notify_event(self, event):
        old_x = self.mousex
        old_y = self.mousey
        self.mousex, self.mousey = self.input_transform.transform_point(
            event.x, event.y)
        dif_x = self.mousex - old_x
        dif_y = self.mousey - old_y
        difs = np.array([dif_x, dif_y])

        if self.mode == Mode.kEditing and self.control_point_index != None:
            multispline = self.multisplines[
                self.control_point_index.multispline_index]
            multispline.updates_for_mouse_move(
                self.control_point_index.control_point_index,
                self.control_point_index.spline_index, self.mousex,
                self.mousey, difs)

            multispline.update_lib_spline()
            self.graph.schedule_recalculate(self.multisplines)
        self.queue_draw()

    def do_scroll_event(self, event):

        self.mousex, self.mousey = self.input_transform.transform_point(
            event.x, event.y)

        step_size = self.pxToM(20)  # px

        if event.direction == Gdk.ScrollDirection.UP:
            # zoom out
            scale_by = step_size
        elif event.direction == Gdk.ScrollDirection.DOWN:
            # zoom in
            scale_by = -step_size
        else:
            return

        scale = (self.field.width + scale_by) / self.field.width

        # This restricts the amount it can be scaled.
        if self.zoom_transform.xx <= 0.5:
            scale = max(scale, 1)
        elif self.zoom_transform.xx >= 16:
            scale = min(scale, 1)

        # undo the scaled translation that the old zoom transform did
        x, y = self.invert(self.zoom_transform).transform_point(
            event.x, event.y)

        # move the origin to point
        self.zoom_transform.translate(x, y)

        # scale from new origin
        self.zoom_transform.scale(scale, scale)

        # move back
        self.zoom_transform.translate(-x, -y)

        self.queue_draw()