# This Python file uses the following encoding: utf-8
from pyqtgraph.Qt import QtCore, QtGui
import pyqtgraph as pg
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import Normalize




class pColorMesh(pg.GraphicsObject):


    def __init__(self, x, y, z, cmap=None):
        pg.GraphicsObject.__init__(self)
        
        self.x = x
        self.y = y
        self.z = z

        if cmap is None:
            self.cmap = plt.cm.viridis

        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(pg.mkPen('w'))
        
        xfn = self.z.shape[0]
        yfn = self.z.shape[1]

        norm = Normalize(vmin=self.z.min(), vmax=self.z.max())
        for xi in range(xfn):
            for yi in range(yfn):
                
                p.drawConvexPolygon(QtCore.QPointF(self.x[xi][yi],     self.y[xi][yi]),
                                    QtCore.QPointF(self.x[xi+1][yi],   self.y[xi+1][yi]),
                                    QtCore.QPointF(self.x[xi+1][yi+1], self.y[xi+1][yi+1]),
                                    QtCore.QPointF(self.x[xi][yi+1],   self.y[xi][yi+1]))

                c = self.cmap(norm(self.z[xi][yi]))[:-1]
                p.setBrush(QtGui.QColor(c[0]*255, c[1]*255, c[2]*255))
        p.end()
    
    def paint(self, p, *args):
        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())


# The example

x = np.array([[1,1,1,1],
              [2,2,2,2],
              [3,3,3,3],
              [4,4,4,4],
              [5,5,5,5]])
y = np.array([[4,8,12,16],
              [2,4,6,8],
              [3,6,9,12],
              [5,10,15,20],
              [6,12,18,24]])
z = np.array([[1,2,3],
              [5,6,7],
              [9,10,11],
              [13,14,15]])

app = QtGui.QApplication([])
mw = QtGui.QMainWindow()
view = pg.GraphicsLayoutWidget()
mw.setCentralWidget(view)
mw.show()
mw.setWindowTitle('pyqtgraph example: pColorMesh')

v = view.addPlot()


item = pColorMesh(x, y, z)
v.addItem(item)

## Start Qt event loop unless running in interactive mode or using pyside.
if __name__ == '__main__':
    import sys
    if (sys.flags.interactive != 1) or not hasattr(QtCore, 'PYQT_VERSION'):
        QtGui.QApplication.instance().exec_()