clusterview2/main_window.py

from functools import partial

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

from clusterview2.exceptions import handle_exceptions
from clusterview2.colors import Color
from clusterview2.mode import Mode
from clusterview2.ui.mode_handlers import (MODE_HANDLER_MAP,
                                           ogl_keypress_handler,
                                           refresh_point_list,
                                           reset_colors,
                                           generate_random_points)
from clusterview2.ui.opengl_widget import (clear_selection, initialize_gl,
                                           mouse_leave, paint_gl, resize_gl,
                                           set_drawing_context)
from clusterview2.points import PointSet
from clusterview2.point_manager import PointManager
from clusterview2.ui.point_list_widget import item_click_handler
from clusterview2_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 clusters 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_clusters.setMinimum(0)
        self.number_of_clusters.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.clustering_button.clicked.connect(self._clustering)
        self.number_of_clusters.valueChanged.connect(self._clustering_enabled)

        self.reset_button.clicked.connect(self._reset)

        # -----------------------------------------------
        # 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

        # Clustering flag so it does not continue to run
        self.clustering_solved = False

    # -----------------------------------------------------------------
    # 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 _clustering(self):
        if len(list(PointManager.point_set.points)) == 0:
            error_dialog = QErrorMessage()
            error_dialog.showMessage('Place points before clustering.')
            error_dialog.exec_()
            return

        clear_selection()
        self._mode = Mode.CLUSTERING
        self.opengl_widget.setCursor(QCursor(Qt.CursorShape.ArrowCursor))
        self.status_bar.showMessage('CLUSTERING')
        self.opengl_widget.update()

    def _reset(self):
        self._off_mode()
        self.number_of_clusters.setEnabled(True)
        self.number_of_clusters.setValue(0)
        self.clustering_button.setEnabled(False)
        self.clustering_solved = False
        PointManager.clusters = []

        for point in PointManager.point_set.points:
            point.weight = 1.0

        reset_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)

    def _clustering_enabled(self):
        self.clustering_button.setEnabled(self.number_of_clusters.value() > 0)

    @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)