from ..Qt import QtGui, QtCore
import math
import numpy as np
from ..colormap import ColorMap
from .GraphicsObject import GraphicsObject
from .. import mkBrush, mkPen
from .. import functions as fn


class NonUniformImage(GraphicsObject):
    """
    **Bases:** :class:`GraphicsObject <pyqtgraph.GraphicsObject>`

    GraphicsObject displaying an image with non-uniform sample points. It's
    commonly used to display 2-d or slices of higher dimensional data that
    have a regular but non-uniform grid e.g. measurements or simulation results.
    """

    def __init__(self, x, y, z, border=None):

        GraphicsObject.__init__(self)

        # convert to numpy arrays
        x = np.asarray(x, dtype=np.float64)
        y = np.asarray(y, dtype=np.float64)
        z = np.asarray(z, dtype=np.float64)

        if x.ndim != 1 or y.ndim != 1:
            raise Exception("x and y must be 1-d arrays.")

        if np.any(np.diff(x) < 0) or np.any(np.diff(y) < 0):
            raise Exception("The values in x and y must be monotonically increasing.")

        if len(z.shape) != 2 or z.shape != (x.size, y.size):
            raise Exception("The length of x and y must match the shape of z.")

        # default colormap (black - white)
        self.cmap = ColorMap(pos=[0.0, 1.0], color=[(0.0, 0.0, 0.0, 1.0), (1.0, 1.0, 1.0, 1.0)])

        self.data = (x, y, z)
        self.lut = None
        self.border = border
        self.generatePicture()

    def setLookupTable(self, lut, autoLevel=False):
        lut.sigLevelsChanged.connect(self.generatePicture)
        lut.gradient.sigGradientChanged.connect(self.generatePicture)
        self.lut = lut

        if autoLevel:
            _, _, z = self.data
            f = z[np.isfinite(z)]
            lut.setLevels(f.min(), f.max())

        self.generatePicture()

    def setColorMap(self, cmap):
        self.cmap = cmap
        self.generatePicture()

    def getHistogram(self, **kwds):
        """Returns x and y arrays containing the histogram values for the current image.
        For an explanation of the return format, see numpy.histogram().
        """

        z = self.data[2]
        z = z[np.isfinite(z)]
        hist = np.histogram(z, **kwds)

        return hist[1][:-1], hist[0]

    def generatePicture(self):

        x, y, z = self.data

        self.picture = QtGui.QPicture()
        p = QtGui.QPainter(self.picture)
        p.setPen(mkPen(None))

        # normalize
        if self.lut is not None:
            mn, mx = self.lut.getLevels()
        else:
            f = z[np.isfinite(z)]
            mn = f.min()
            mx = f.max()

        # draw the tiles
        for i in range(x.size):
            for j in range(y.size):

                value = z[i, j]

                if np.isneginf(value):
                    value = 0.0
                elif np.isposinf(value):
                    value = 1.0
                elif math.isnan(value):
                    continue  # ignore NaN
                else:
                    value = (value - mn) / (mx - mn)  # normalize

                if self.lut:
                    color = self.lut.gradient.getColor(value)
                else:
                    color = self.cmap.mapToQColor(value)

                p.setBrush(mkBrush(color))

                # left, right, bottom, top
                l = x[0] if i == 0 else (x[i - 1] + x[i]) / 2
                r = (x[i] + x[i + 1]) / 2 if i < x.size - 1 else x[-1]
                b = y[0] if j == 0 else (y[j - 1] + y[j]) / 2
                t = (y[j] + y[j + 1]) / 2 if j < y.size - 1 else y[-1]

                p.drawRect(QtCore.QRectF(l, t, r - l, b - t))

        if self.border is not None:
            p.setPen(self.border)
            p.setBrush(fn.mkBrush(None))
            p.drawRect(self.boundingRect())

        p.end()

        self.update()

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

    def boundingRect(self):
        return QtCore.QRectF(self.picture.boundingRect())