clusterview/main_window.py

from functools import partial

from PyQt5.QtCore import Qt
from PyQt5.QtGui import QCursor
from PyQt5.QtWidgets import QFileDialog, QInputDialog, QMainWindow

from clusterview.exceptions import handle_exceptions
from clusterview.colors import Color
from clusterview.mode import Mode
from clusterview.mode_handlers import (group, MODE_HANDLER_MAP,
                                       ogl_keypress_handler,
                                       refresh_point_list,
                                       reset_centroid_count_and_colors,
                                       generate_random_points)
from clusterview.opengl_widget import (clear_selection, initialize_gl,
                                       mouse_leave, paint_gl, resize_gl,
                                       set_drawing_context)
from clusterview.points import PointSet
from clusterview.point_manager import PointManager
from clusterview.point_list_widget import item_click_handler
from clusterview_ui import Ui_MainWindow


class MainWindow(QMainWindow, Ui_MainWindow):
    """
    A wrapper class for handling creating a window based
    on the `clusterview_ui.py` code generated from
    `clusterview.ui`.
    """

    # This is a static mode variable since there will only ever
    # be one MainWindow.
    __mode = Mode.OFF

    def __init__(self, parent=None):
        super(MainWindow, self).__init__(parent)
        self.setupUi(self)

        # Size of point for drawing
        self.__point_size = 8

        # TODO: THESE ARE HARD CODED TO THE CURRENT QT WIDGET SIZES
        #       FIX THIS PROPERLY WITH A RESIZE EVENT DETECT.
        # PointManager is a class that is filled with static methods
        # designed for managing state.
        self.__viewport_width = 833
        self.__viewport_height = 656

        PointManager.point_set = PointSet(self.__point_size,
                                          self.__viewport_width,
                                          self.__viewport_height)

        # Spin box should only allow the number of centroids to be no
        # greater than the number of supported colors minus 2 to exclude
        # the color for selection (Color.BLUE) and the default color for points
        # (Color.GREY).
        self.number_of_centroids.setMinimum(0)
        self.number_of_centroids.setMaximum(Color.count() - 2)

        # We only need to set the context in our OpenGL state machine
        # wrapper once here since the window is fixed size.
        # If we allow resizing of the window, the context must be updated
        # each resize so that coordinates are converted from screen (x, y)
        # to OpenGL coordinates properly.
        set_drawing_context(self)

        # Enables mouse tracking on the viewport so mouseMoveEvents are
        # tracked and fired properly.
        self.opengl_widget.setMouseTracking(True)

        # Enable keyboard input capture on the OpenGL Widget
        self.opengl_widget.setFocusPolicy(Qt.StrongFocus)

        # Here we partially apply the key press handler with self to
        # create a new function that only expects the event `keyPressEvent`
        # expects. In this way, we've snuck the state of the opengl_widget
        # into the function so that we can modify it as we please.
        self.opengl_widget.keyPressEvent = partial(ogl_keypress_handler, self)

        # Same story here but this time with the itemClicked event
        # so that when an element is clicked on in the point list it will
        # highlight.
        self.point_list_widget.itemClicked.connect(partial(item_click_handler,
                                                           self))

        self.choose_centroids_button.clicked.connect(self.__choose_centroids)

        self.group_button.clicked.connect(self.__group)

        self.reset_button.clicked.connect(self.__reset_grouping)

        # -----------------------------------------------
        # OpenGL Graphics Handlers are set
        # here and defined in clusterview.opengl_widget.
        # -----------------------------------------------
        self.opengl_widget.initializeGL = initialize_gl
        self.opengl_widget.paintGL = paint_gl
        self.opengl_widget.resizeGL = resize_gl
        self.opengl_widget.leaveEvent = partial(mouse_leave, self)

        # -------------------------------------
        # UI Handlers
        # -------------------------------------
        self.action_add_points.triggered.connect(self.__add_points)
        self.action_edit_points.triggered.connect(self.__edit_points)
        self.action_delete_points.triggered.connect(self.__delete_points)
        self.action_move_points.triggered.connect(self.__move_points)

        (self.action_generate_random_points
            .triggered.connect(self.__generate_random_points))

        self.action_save_point_configuration.triggered.connect(
            self.__save_points_file)

        self.action_load_point_configuration.triggered.connect(
            self.__open_points_file)

        self.action_exit.triggered.connect(self.__close_event)

        # Override handler for mouse press so we can draw points based on
        # the OpenGL coordinate system inside of the OpenGL Widget.
        self.opengl_widget.mousePressEvent = self.__ogl_click_dispatcher
        self.opengl_widget.mouseMoveEvent = self.__ogl_click_dispatcher
        self.opengl_widget.mouseReleaseEvent = self.__ogl_click_dispatcher

    # -----------------------------------------------------------------
    # Mode changers - these will be used to signal the action in the
    #                 OpenGL Widget.
    # -----------------------------------------------------------------
    def __off_mode(self):
        self.__mode = Mode.OFF
        self.opengl_widget.setCursor(QCursor(Qt.CursorShape.ArrowCursor))
        self.status_bar.showMessage("")
        clear_selection()
        self.opengl_widget.update()

    def __add_points(self):
        self.__mode = Mode.ADD
        self.opengl_widget.setCursor(QCursor(Qt.CursorShape.CrossCursor))
        self.status_bar.showMessage("ADD MODE")
        clear_selection()
        self.opengl_widget.update()

    def __edit_points(self):
        self.__mode = Mode.EDIT
        self.opengl_widget.setCursor(QCursor(Qt.CursorShape.CrossCursor))
        self.status_bar.showMessage("EDIT MODE")
        clear_selection()
        self.opengl_widget.update()

    def __delete_points(self):
        self.__mode = Mode.DELETE
        self.opengl_widget.setCursor(QCursor(
            Qt.CursorShape.PointingHandCursor))
        self.status_bar.showMessage("DELETE MODE")
        clear_selection()
        self.opengl_widget.update()

    def __move_points(self):
        self.__mode = Mode.MOVE
        self.opengl_widget.setCursor(QCursor(Qt.CursorShape.SizeAllCursor))
        self.status_bar.showMessage("MOVE MODE - PRESS ESC OR SWITCH MODES" +
                                    "TO CANCEL SELECTION")
        clear_selection()
        self.opengl_widget.update()

    def __choose_centroids(self):
        self.__mode = Mode.CHOOSE_CENTROIDS
        self.opengl_widget.setCursor(QCursor(Qt.CursorShape.CrossCursor))
        self.status_bar.showMessage("CHOOSE CENTROIDS")

        clear_selection()
        self.opengl_widget.update()

    def __group(self):
        self.__mode = Mode.GROUP

        self.opengl_widget.setCursor(QCursor(Qt.CursorShape.ArrowCursor))
        self.status_bar.showMessage("GROUPING")
        clear_selection()
        group(self)
        self.__off_mode()
        self.opengl_widget.update()

    def __reset_grouping(self):
        self.__off_mode()
        self.number_of_centroids.setEnabled(True)
        self.number_of_centroids.setValue(0)
        self.choose_centroids_button.setEnabled(True)
        self.solve_button.setEnabled(False)
        self.group_button.setEnabled(False)
        PointManager.centroids = []
        reset_centroid_count_and_colors()

    def __generate_random_points(self):
        value, ok = QInputDialog.getInt(self, "Number of Points",
                                        "Number of Points:", 30, 30, 3000, 1)

        if ok:
            self.__mode = Mode.ADD
            generate_random_points(value,
                                   (self.__viewport_width - self.__point_size),
                                   (self.__viewport_height - self.__point_size)
                                   )
            self.__mode = Mode.OFF

            self.opengl_widget.update()

            refresh_point_list(self)

    @property
    def mode(self):
        """
        Function designed to be used from a context
        to get the current mode.
        """
        return self.__mode

    @mode.setter
    def mode(self, mode):
        self.__mode = mode

    def __close_event(self, event):
        import sys
        sys.exit(0)

    def __open_points_file(self):
        ofile, _ = QFileDialog.getOpenFileName(self,
                                               "Open Point Configuration",
                                               "",
                                               "JSON files (*.json)")
        if ofile:
            self.__mode = Mode.LOADED

            PointManager.load(ofile)

            self.opengl_widget.update()

            refresh_point_list(self)

    def __save_points_file(self):
        file_name, _ = (QFileDialog.
                        getSaveFileName(self,
                                        "Save Point Configuration",
                                        "",
                                        "JSON Files (*.json)"))
        if file_name:
            PointManager.save(file_name)

    @handle_exceptions
    def __ogl_click_dispatcher(self, event):
        """
        Mode dispatcher for click actions on the OpenGL widget.
        """
        # Map from Mode -> function
        # where the function is a handler for the
        # OpenGL event. The context passed to these functions allows
        # them to modify on screen widgets such as the QOpenGLWidget and
        # QListWidget.
        MODE_HANDLER_MAP[self.__mode](self, event)