clusterview2/clusterview2/ui/opengl_widget.py

"""
This module defines functions that need to be overwritten
in order for OpenGL to work with the main window. This
module is named the same as the actual widget in order
to make namespacing consistent.

To be clear, the actual widget is defined in the UI
generated code - `clusterview_ui.py`. The functions
here are imported as overrides to the OpenGL functions of
that widget.

It should be split up into a few more separate files eventually...
Probably even into it's own module folder.
"""

from OpenGL.GL import (glBegin, glClearColor, glColor3f,
                       glEnd, GL_LINE_LOOP, GL_POINTS,
                       glPointSize, glVertex3f, glViewport)

from clusterview2.colors import Color, COLOR_TO_RGBA
from clusterview2.exceptions import (handle_exceptions,
                                     InvalidStateError)
from clusterview2.mode import Mode
from clusterview2.point_manager import PointManager

# Constants set based on the size of the window.
__BOTTOM_LEFT = (0, 0)
__WIDTH = None
__HEIGHT = None

# State variables for a move selection bounding box.
# There are always reset to None after a selection has been made.
__move_bb_top_left = None
__move_bb_bottom_right = None

# Module-global state variables for our drawing
# state machine.
#
# Below functions have to mark these as `global` so
# the interpreter knows that the variables are not
# function local.
__current_context = None
__current_event = None


# TODO: This should live inside of a class as static methods with the
#       globals moved into the static scope to make this nicer...once you
#       get it running before doing kmeans make this modification.

def set_drawing_context(ctx):
    """
    Sets the drawing context so that drawing functions can properly
    interact with the widget.
    """
    global __current_context

    __current_context = ctx


def set_drawing_event(event):
    """
    State machine event management function.

    @param event The event.
    """
    global __current_context
    global __current_event

    if __current_context is None:
        raise InvalidStateError('Drawing context must be set before setting ' +
                                'drawing mode')

    if event is not None:
        __current_event = event


def mouse_leave(ctx, event):
    """
    The leave event for the OpenGL widget to properly reset the mouse
    position label.

    @param ctx The context.
    @param event The event.
    """
    ctx.mouse_position_label.setText('')


def set_move_bb_top_left(x, y):
    """
    Called to set the move bounding box's top left corner.

    @param x The x-coordinate.
    @param y The y-coordinate.
    """
    global __move_bb_top_left

    __move_bb_top_left = (x, y)


def set_move_bb_bottom_right(x, y):
    """
    Called to set the move bounding box's bottom right corner.

    @param x The x-coordinate.
    @param y The y-coordinate.
    """
    global __move_bb_bottom_right

    __move_bb_bottom_right = (x, y)


def get_bb_top_left():
    return __move_bb_top_left


def get_bb_bottom_right():
    return __move_bb_bottom_right


def reset_move_bbs():
    global __move_bb_top_left
    global __move_bb_bottom_right

    __move_bb_top_left = None
    __move_bb_bottom_right = None


def initialize_gl():
    """
    Initializes the OpenGL context on the Window.
    """

    # Set white background
    glClearColor(255, 255, 255, 0)


def resize_gl(w, h):
    """
    OpenGL resize handler used to get the current viewport size.

    @param w The new width.
    @param h The new height.
    """
    global __WIDTH
    global __HEIGHT

    __WIDTH = __current_context.opengl_widget.width()
    __HEIGHT = __current_context.opengl_widget.height()


def viewport_width():
    return __WIDTH


def viewport_height():
    return __HEIGHT


@handle_exceptions
def paint_gl():
    """
    Stock PaintGL function from OpenGL that switches
    on the current mode to determine what action to
    perform on the current event.
    """
    if(__current_context.mode is Mode.OFF and
       not PointManager.point_set.empty()):

        # We want to redraw on any change to Mode.OFF so points are preserved -
        # without this, any switch to Mode.OFF will cause a blank screen to
        # render.
        draw_points(PointManager.point_set)

    if (__current_context.mode in [Mode.ADD, Mode.EDIT,
                                   Mode.MOVE, Mode.DELETE] and
            __current_event is None and PointManager.point_set.empty()):
        return

    if (__current_context.mode in [Mode.ADD, Mode.EDIT, Mode.DELETE] and
            PointManager.point_set.empty()):
        return

    if (__current_context.mode is Mode.ADD or
            __current_context.mode is Mode.DELETE or
            __current_context.mode is Mode.LOADED or
            __current_context.mode is Mode.CHOOSE_CENTROIDS or
            __current_context.mode is Mode.UNWEIGHTED_CLUSTERING or
            __current_context.mode is Mode.WEIGHTED_CLUSTERING):

        draw_points(PointManager.point_set)

    elif __current_context.mode is Mode.EDIT:
        raise NotImplementedError('Drawing for EDIT not implemented.')

    elif __current_context.mode is Mode.MOVE:
        # We have to repeatedly draw the points while we are showing the
        # move box.
        if not PointManager.point_set.empty():
            draw_points(PointManager.point_set)

            draw_selection_box(Color.BLACK)

        if (__move_bb_top_left is not None and
                __move_bb_bottom_right is not None):

            # Mark points that are selected in the bounding box
            # and draw them using the normal function
            highlight_selection()
            draw_points(PointManager.point_set)


def __clamp_x(x):
    """
    X-coordinate clamping function that goes from mouse coordinates to
    OpenGL coordinates.

    @param x The x-coordinate to clamp.
    @returns The clamped x coordinate.
    """
    x_w = (x / (__WIDTH / 2.0) - 1.0)
    return x_w


def __clamp_y(y):
    """
    Y-coordinate clamping function that goes from mouse coordinates to
    OpenGL coordinates.

    @param y The y-coordinate to clamp.
    @returns The clamped y coordinate.
    """
    y_w = -1.0 * (y / (__HEIGHT / 2.0) - 1.0)
    return y_w


def box_hit(tx, ty, x1, y1, x2, y2):
    """
    Calculates whether or not a given point collides with the given bounding
    box.

    @param tx The target x.
    @param ty The target y.
    @param x1 The top left x.
    @param y1 The top left y.
    @param x2 The bottom left x.
    @param y2 The bottom left y.
    """

    # The box in this case is flipped - the user started at the bottom right
    # corner. Pixel-wise top left is (0, 0) and bottom right is
    # (screen_x, screen_y)
    if x1 > x2 and y1 > y2:
        return (tx <= x1 and
                tx >= x2 and
                ty <= y1 and
                ty >= y2)

    # The box in this case started from the top right
    if x1 > x2 and y1 < y2:
        return (tx <= x1 and
                tx >= x2 and
                ty >= y1 and
                ty <= y2)

    # The box in this case started from the bottom left
    if x1 < x2 and y1 > y2:
        return (tx >= x1 and
                tx <= x2 and
                ty <= y1 and
                ty >= y2)

    # Normal condition: Box starts from the top left
    return (tx >= x1 and
            tx <= x2 and
            ty >= y1 and
            ty <= y2)


def highlight_selection():
    """
    Given the current move bounding box, highlights any points inside it.
    """

    top_left = get_bb_top_left()
    bottom_right = get_bb_bottom_right()

    for point in PointManager.point_set.points:
        if box_hit(point.x, point.y, top_left[0], top_left[1],
                   bottom_right[0], bottom_right[1]):

            point.select()
        else:
            point.unselect()


def draw_selection_box(color):
    """
    When the move bounding box state is populated and the mode is set
    to MODE.Move this function will draw the selection bounding box.

    @param color The color Enum.
    """
    global __current_context

    if __current_context is None:
        raise InvalidStateError('Drawing context must be set before setting ' +
                                'drawing mode')

    if not isinstance(color, Color):
        raise ValueError('Color must exist in the Color enumeration')

    if __move_bb_top_left is None or __move_bb_bottom_right is None:
        # Nothing to draw.
        return

    ct = COLOR_TO_RGBA[color]

    glViewport(0, 0, __WIDTH, __HEIGHT)

    # Top right corner has the same x as the bottom right
    # and same y as the top left.
    top_right_corner = (__move_bb_bottom_right[0], __move_bb_top_left[1])

    # Bottom left corner has the same x as the top left and
    # same y as the bottom right.
    bottom_left_corner = (__move_bb_top_left[0], __move_bb_bottom_right[1])

    glBegin(GL_LINE_LOOP)
    glColor3f(ct[0], ct[1], ct[2])

    glVertex3f(__clamp_x(__move_bb_top_left[0]),
               __clamp_y(__move_bb_top_left[1]),
               0.0)

    glVertex3f(__clamp_x(top_right_corner[0]),
               __clamp_y(top_right_corner[1]),
               0.0)

    glVertex3f(__clamp_x(__move_bb_bottom_right[0]),
               __clamp_y(__move_bb_bottom_right[1]),
               0.0)

    glVertex3f(__clamp_x(bottom_left_corner[0]),
               __clamp_y(bottom_left_corner[1]),
               0.0)

    glEnd()


def clear_selection():
    """
    A helper designed to be called from the main window
    in order to clear the selection internal to the graphics
    and mode files. This way you dont have to do something
    before the selection clears.
    """
    if not PointManager.point_set.empty():
        PointManager.point_set.clear_selection()


def draw_points(point_set):
    """
    Simple point drawing function.

    Given a coordinate (x, y), and a Color enum this
    function will draw the given point with the given
    color.

    @param point_set The PointSet to draw.
    @param color The Color Enum.
    """
    global __current_context

    if __current_context is None:
        raise InvalidStateError('Drawing context must be set before setting ' +
                                'drawing mode')

    glViewport(0, 0, __WIDTH, __HEIGHT)
    glPointSize(PointManager.point_set.point_size)

    glBegin(GL_POINTS)
    for point in point_set.points:

        if point.selected:
            blue = COLOR_TO_RGBA[Color.BLUE]
            glColor3f(blue[0], blue[1], blue[2])
        else:
            ct = COLOR_TO_RGBA[point.color]
            glColor3f(ct[0], ct[1], ct[2])

        glVertex3f(__clamp_x(point.x),
                   __clamp_y(point.y),
                   0.0)                 # Z is currently fixed to 0
    glEnd()