nmreval/src/gui_qt/graphs/draw_objects.py

from __future__ import annotations

import uuid

from pyqtgraph import TextItem, mkPen, GraphicsObject, mkBrush
from numpy import log10

from gui_qt.Qt import QtCore, QtGui
from gui_qt.lib.pg_objects import LogInfiniteLine, PlotItem


__all__ = ['LineObject', 'MultipointObject', 'RectangleObject', 'TextObject', 'EllipseObject']


class BaseObject:
    def __init__(self, color=None):
        self.id = str(uuid.uuid4())
        self.color = color

        self.drawing = None

    def __str__(self):
        raise NotImplementedError


class LineObject(BaseObject):
    def __init__(self, pos: float, angle: float, **kwargs):
        super().__init__(**kwargs)

        self.pos = pos
        self.angle = angle

        self.drawing = LogInfiniteLine(pos=self.pos, angle=self.angle, pen=mkPen(color=self.color.rgb()))

    def __str__(self):
        return f'{"x" if self.angle==90 else "y"}={self.pos}'


class MultipointObject(BaseObject):
    def __init__(self, pts: list, closed: bool = False, **kwargs):
        super().__init__(**kwargs)
        x, y = zip(*pts)
        self.closed = closed

        if self.closed:
            x += (x[0],)
            y += (y[0],)
        self._x = x
        self._y = y

        self.drawing = PlotItem(x=self._x, y=self._y, pen=mkPen(color=self.color.rgb()))

    def __str__(self):
        return f'{len(self._y)-int(self.closed)}-pts'


class Rectangle(GraphicsObject):
    # adapted from pyqtgraph example on custom objects

    def __init__(self, left: tuple, right: tuple):
        GraphicsObject.__init__(self)

        self.p0 = left
        self.delta = [right[0]-self.p0[0], right[1]-self.p0[1]]

        self.x_disp = self.p0[0]
        self.dx = self.delta[0]
        self.y_disp = self.p0[1]
        self.dy = self.delta[1]

        self.log_mode = [False, False]
        self.generatePicture()

    def generatePicture(self):
        ## pre-computing a QPicture object allows paint() to run much more quickly,
        ## rather than re-drawing the shapes every time.
        self.picture = QtGui.QPicture()
        p = QtGui.QPainter(self.picture)
        p.setPen(mkPen('w'))
        p.setBrush(mkBrush('r'))
        p.drawRect(QtCore.QRectF(self.x_disp, self.y_disp, self.dx, self.dy))
        p.end()

    def paint(self, p, *_):
        p.drawPicture(0, 0, self.picture)

    def boundingRect(self):
        ## boundingRect _must_ indicate the entire area that will be drawn on
        ## or else we will get artifacts and possibly crashing.
        ## (in this case, QPicture does all the work of computing the bouning rect for us)
        return QtCore.QRectF(self.picture.boundingRect())

    def setLogMode(self, x_mode, y_mode):
        if self.log_mode == [x_mode, y_mode]:
            return

        if self.log_mode[0] != x_mode:
            if x_mode:
                x0 = self.p0[0]
                xmax = log10(x0+self.delta[0])
                self.x_disp = log10(x0)
                self.dx = xmax - self.x_disp
            else:
                self.x_disp = self.p0[0]
                self.dx = self.delta[0]

        if self.log_mode[1] != y_mode:
            if y_mode:
                y0 = self.p0[1]
                ymax = log10(y0 + self.delta[1])
                self.y_disp = log10(y0)
                self.dy = ymax - self.y_disp
            else:
                self.y_disp = self.p0[1]
                self.dy = self.delta[1]

        self.generatePicture()

        self.log_mode = [x_mode, y_mode]

        self.informViewBoundsChanged()

class RectangleObject(BaseObject):
    def __init__(self, left: tuple, right: tuple, **kwargs):
        super().__init__(**kwargs)

        self.left = left
        self.right = right

        self.drawing = Rectangle(left, right)

    def __str__(self):
        return f'Rectangle {self.left} to {self.right}'


class TextObject(BaseObject):
    def __init__(self, text, pos, **kwargs):
        super().__init__(**kwargs)

        self.pos = pos
        self.drawing = TextItem(text, anchor=(0.5, 0.5), color=self.color.rgb())
        self.drawing.setPos(*pos)

    def __str__(self):
        return f'{self.drawing.toPlainText()} at {self.pos}'


class EllipseObject(BaseObject):
    def __init__(self, center: tuple, axes: tuple, **kwargs):
        super().__init__(**kwargs)

        self.center = center
        self.axes = axes

        self.drawing = Ellipse(center, axes)

    def __str__(self):
        return f'Ellipse at {self.center}'

class Ellipse(GraphicsObject):
    # adapted from pyqtgraph example on custom objects

    def __init__(self, center: tuple, axes: tuple):
        GraphicsObject.__init__(self)

        self.center = center
        self.ax = axes

        self.center_disp = self.center
        self.ax_disp = self.ax

        self.log_mode = [False, False]
        self.generatePicture()

    def generatePicture(self):
        ## pre-computing a QPicture object allows paint() to run much more quickly,
        ## rather than re-drawing the shapes every time.
        self.picture = QtGui.QPicture()
        p = QtGui.QPainter(self.picture)
        p.setPen(mkPen('w'))
        p.setBrush(mkBrush('r'))
        p.drawEllipse(QtCore.QPointF(*self.center_disp), *self.ax_disp)
        p.end()

    def paint(self, p, *_):
        p.drawPicture(0, 0, self.picture)

    def boundingRect(self):
        ## boundingRect _must_ indicate the entire area that will be drawn on
        ## or else we will get artifacts and possibly crashing.
        ## (in this case, QPicture does all the work of computing the bouning rect for us)
        return QtCore.QRectF(self.picture.boundingRect())

    def setLogMode(self, x_mode, y_mode):
        if self.log_mode == [x_mode, y_mode]:
            return

        if self.log_mode[0] != x_mode:
            if x_mode:
                x0 = self.p0[0]
                xmax = log10(x0+self.delta[0])
                self.x_disp = log10(x0)
                self.dx = xmax - self.x_disp
            else:
                self.x_disp = self.p0[0]
                self.dx = self.delta[0]

        if self.log_mode[1] != y_mode:
            if y_mode:
                y0 = self.p0[1]
                ymax = log10(y0 + self.delta[1])
                self.y_disp = log10(y0)
                self.dy = ymax - self.y_disp
            else:
                self.y_disp = self.p0[1]
                self.dy = self.delta[1]

        self.generatePicture()

        self.log_mode = [x_mode, y_mode]

        self.informViewBoundsChanged()

    def get_state(self):
        return {'center': self.center, 'axes': self.ax}