import wx
import wx.glcanvas
import numpy

def makeWindow():
    global window
    window = ViewerWindow(None, tileShape = (3, 3), tileSize = 512)

    # Generate a test image to display -- a sequence of squares.
    width = 2056
    height = 2548
    data = numpy.zeros((width, height), dtype = numpy.uint16)
    base = 0
    for i in xrange(0, width, 512):
        for j in xrange(0, height, 512):
            data[i:i+512, j:j+512] = base
            base += 50
    # Downsample to make the image reasonably displayable on a normal monitor.
    window.viewer.setImage(data[::2, ::2])

class App(wx.App):
    def OnInit(self):
        makeWindow()
        return True

from OpenGL.GL import *

## Simple window that contains a GLViewer instance.
class ViewerWindow(wx.Frame):
    def __init__(self, parent, tileShape, tileSize, *args, **kwargs):
        wx.Frame.__init__(self, parent, *args, **kwargs)

        self.viewer = GLViewer(self, tileShape, tileSize, **kwargs)
        self.Show()
        self.Bind(wx.EVT_SIZE, self.onSize)

    def onSize(self, event):
        self.viewer.setSize(list(self.GetSize()))
        event.Skip()


class GLViewer(wx.glcanvas.GLCanvas):
    ## Instantiate.
    def __init__(self, parent, tileShape, tileSize, *args, **kwargs):
        wx.glcanvas.GLCanvas.__init__(self, parent, *args, **kwargs)

        ## Edge length of one tile.
        self.tileSize = tileSize

        ## Shape of our tile grid.
        self.tileShape = tileShape

        ## 2D array of Image instances. We'll create these when we get
        # some image data to work with.
        self.tiles = []

        ## Whether or not we've done some one-time initialization work.
        self.haveInitedGL = False

        self.scaleX = self.scaleY = 1.0

        wx.EVT_PAINT(self, self.OnPaint)


    def InitGL(self):
        self.w, self.h = self.GetClientSizeTuple()
        self.SetCurrent()
        glClearColor(0.3, 0.3, 0.3, 0.0)   ## background color

        self.haveInitedGL = True


    ## Receive a new image and parcel it up into our smaller Image instances.
    # Create those instances at this time if we don't have them already.
    def setImage(self, newImage):
        self.SetCurrent()

        print "Updating new image to one of shape",newImage.shape

        # Ensure that we have the right amount of data to put into our
        # Images, by filling in any missing data with zeros. There's an
        # implicit assumption here that newImage is smaller than or
        # equal in size to our tiles' total combined area.
        width = self.tileShape[0] * self.tileSize
        height = self.tileShape[1] * self.tileSize
        paddedImage = numpy.zeros((width, height), dtype = newImage.dtype)
        paddedImage[:newImage.shape[0], :newImage.shape[1]] = newImage

        print "Made padded image of shape",paddedImage.shape

        for i in xrange(self.tileShape[0]):
            self.tiles.append([])
            for j in xrange(self.tileShape[1]):
                subData = paddedImage[
                        i * self.tileSize : (i + 1) * self.tileSize,
                        j * self.tileSize : (j + 1) * self.tileSize]
                self.tiles[i].append(Image(subData))

        wx.CallAfter(self.changeHistScale, newImage.min(), newImage.max())

        self.Refresh()


    ## Rescale our images so smin displays as black and smax as white.
    def changeHistScale(self, smin ,smax):
        for i in xrange(self.tileShape[0]):
            for j in xrange(self.tileShape[1]):
                self.tiles[i][j].refresh(smin, smax)
        self.Refresh(False)


    def setSize(self, size):
        xSize = float(self.tileShape[0] * self.tileSize)
        ySize = float(self.tileShape[1] * self.tileSize)
        self.scaleX = size[0] / xSize
        self.scaleY = size[1] / ySize
        self.w, self.h = size
        self.Refresh(0)


    def OnPaint(self, event):
        try:
            dc = wx.PaintDC(self)
        except:
            return

        if not self.haveInitedGL:
            self.InitGL()

        self.SetCurrent()

        glViewport(0, 0, self.w, self.h)
        glMatrixMode (GL_PROJECTION)
        glLoadIdentity ()
        glOrtho (0, self.w, 0, self.h, 1., -1.)
        glMatrixMode (GL_MODELVIEW)
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)

        glPushMatrix()
        glLoadIdentity()

        glScalef(self.scaleX, self.scaleY, 1)

        glEnable(GL_TEXTURE_2D)

        for i in xrange(self.tileShape[0]):
            for j in xrange(self.tileShape[1]):
                glPushMatrix()
                glTranslatef(i * self.tileSize, j * self.tileSize, 0)
                self.tiles[i][j].render()
                glPopMatrix()

        glDisable(GL_TEXTURE_2D)
        glPopMatrix()

        glFlush()
        self.SwapBuffers()


import OpenGL.GL as GL

## This class handles display of a single 2D array of pixel data.
class Image:
    def __init__(self, imageData):
        self.imageData = imageData
        self.bindTexture()


    def bindTexture(self):
        pic_ny, pic_nx = self.imageData.shape

        # Generate texture sizes that are powers of 2
        tex_nx = 2
        while tex_nx < pic_nx:
            tex_nx *= 2
        tex_ny = 2
        while tex_ny < pic_ny:
            tex_ny *= 2

        self.picTexRatio_x = float(pic_nx) / tex_nx
        self.picTexRatio_y = float(pic_ny) / tex_ny

        self.textureID = GL.glGenTextures(1)
        GL.glBindTexture(GL.GL_TEXTURE_2D, self.textureID)

        # Define this new texture object based on self.imageData's geometry
        GL.glTexParameteri(GL.GL_TEXTURE_2D, GL.GL_TEXTURE_MIN_FILTER,
                GL.GL_NEAREST)
        GL.glTexParameteri(GL.GL_TEXTURE_2D, GL.GL_TEXTURE_MAG_FILTER,
                GL.GL_NEAREST)

        imgType = self.imageData.dtype.type
        GL.glTexImage2D(GL.GL_TEXTURE_2D,0,  GL.GL_RGB, tex_nx,tex_ny, 0,
                     GL.GL_LUMINANCE, GL.GL_UNSIGNED_SHORT, None)


    # Re-load our image data onto the card, e.g. because our blackpoint and
    # whitepoint have changed.
    def refresh(self, imageMin, imageMax):
        minMaxRange = float(imageMax - imageMin)
        if abs(imageMax - imageMin) < 1:
            minMaxRange = 1

        imgType = self.imageData.dtype.type
        fBias = -imageMin / minMaxRange
        f = ((1 << 16) - 1) / minMaxRange

        GL.glBindTexture(GL.GL_TEXTURE_2D, self.textureID)

        GL.glPixelTransferf(GL.GL_RED_SCALE,   f)
        GL.glPixelTransferf(GL.GL_GREEN_SCALE, f)
        GL.glPixelTransferf(GL.GL_BLUE_SCALE,  f)

        GL.glPixelTransferf(GL.GL_RED_BIAS,   fBias)
        GL.glPixelTransferf(GL.GL_GREEN_BIAS, fBias)
        GL.glPixelTransferf(GL.GL_BLUE_BIAS,  fBias)

        GL.glPixelTransferf(GL.GL_MAP_COLOR, False)

        GL.glPixelStorei(GL.GL_UNPACK_SWAP_BYTES,
                not self.imageData.dtype.isnative)
        GL.glPixelStorei(GL.GL_UNPACK_ALIGNMENT, self.imageData.itemsize)

        imgString = self.imageData.tostring()

        pic_ny, pic_nx = self.imageData.shape

        GL.glTexSubImage2D(GL.GL_TEXTURE_2D, 0, 0, 0, pic_nx, pic_ny,
                GL.GL_LUMINANCE, GL.GL_UNSIGNED_SHORT, imgString)


    def render(self):
        cx,cy = self.imageData.shape[-1]/2., self.imageData.shape[-2]/2.
        GL.glBindTexture(GL.GL_TEXTURE_2D, self.textureID)

        GL.glBegin(GL.GL_QUADS)

        pic_ny, pic_nx = self.imageData.shape

        ###//(0,0) at left bottom
        GL.glTexCoord2f(0, 0)
        GL.glVertex2i(0, 0)

        GL.glTexCoord2f(self.picTexRatio_x, 0)
        GL.glVertex2i(pic_nx, 0)

        GL.glTexCoord2f(self.picTexRatio_x, self.picTexRatio_y)
        GL.glVertex2i(pic_nx, pic_ny)

        GL.glTexCoord2f(0, self.picTexRatio_y)
        GL.glVertex2i(0, pic_ny)

        GL.glEnd()


app = App(redirect = False)
app.MainLoop()