import numpy
from OpenGL.GL import *
import threading
import wx
import wx.glcanvas

import numpy
import OpenGL.GL as GL
import threading

## Maps numpy datatypes to OpenGL datatypes
dtypeToGlTypeMap = {
    numpy.uint8: GL.GL_UNSIGNED_BYTE,
    numpy.uint16: GL.GL_UNSIGNED_SHORT,
    numpy.int16: GL.GL_SHORT,
    numpy.float32: GL.GL_FLOAT,
    numpy.float64: GL.GL_FLOAT,
    numpy.int32: GL.GL_FLOAT,
    numpy.uint32: GL.GL_FLOAT,
    numpy.complex64: GL.GL_FLOAT,
    numpy.complex128: GL.GL_FLOAT,
}

## Maps numpy datatypes to the maximum value the datatype can represent
dtypeToMaxValMap = {
    numpy.uint16: (1 << 16) - 1,
    numpy.int16: (1 << 15) - 1,
    numpy.uint8: (1 << 8) - 1, 
    numpy.bool_: (1 << 8) - 1,
    numpy.float32: 1
}

## This class handles display of a single 2D array of pixel data.
class Image:
    def __init__(self):
        self.imageData = None
        self.imageMin = None
        self.imageMax = None
        self.textureID = None
        self.color = (1, 1, 1)
        self.lock = threading.Lock()
        
        self.bindTexture()
        self.refresh()

    def bindTexture(self):
        if self.imageData is None:
            return
        self.lock.acquire()
        pic_ny, pic_nx = self.imageData.shape

        if self.imageMin == 0 and self.imageMax == 0:
            self.imageMin = self.imageData.min()
            self.imageMax = self.imageData.max()

        # 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
        if imgType not in dtypeToGlTypeMap:
            raise ValueError, "Unsupported data mode %s" % str(imgType)
        GL.glTexImage2D(GL.GL_TEXTURE_2D,0,  GL.GL_RGB, tex_nx,tex_ny, 0,
                     GL.GL_LUMINANCE, dtypeToGlTypeMap[imgType], None)
        self.lock.release()

    def refresh(self):
        if self.imageData is None:
            return
        self.bindTexture()
        minMaxRange = float(self.imageMax - self.imageMin)
        if abs(self.imageMax - self.imageMin) < 1:
            minMaxRange = 1

        imgType = self.imageData.dtype.type
        fBias = -self.imageMin / minMaxRange
        f = dtypeToMaxValMap[imgType] / 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

        if imgType not in dtypeToGlTypeMap:
            raise ValueError, "Unsupported data mode %s" % str(imgType)
        GL.glTexSubImage2D(GL.GL_TEXTURE_2D, 0, 0, 0, pic_nx, pic_ny,
                GL.GL_LUMINANCE, dtypeToGlTypeMap[imgType], imgString)


    def render(self):
        if self.imageData is None or self.textureID is None:
            return
        self.refresh()
        
        GL.glPushMatrix()

        GL.glColor3fv(self.color)
        
        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()
        GL.glPopMatrix()


    ## Free the allocated GL texture
    def wipe(self):
        self.lock.acquire()
        if self.textureID is not None:
            GL.glDeleteTextures(self.textureID)
        self.textureID = None
        self.lock.release()
   

    ## Accept a new array of image data.
    def updateImage(self, imageData, imageMin, imageMax):
        self.imageData = imageData
        self.imageMin = imageMin
        self.imageMax = imageMax

        self.wipe()
        self.refresh()



colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255)]


## Simple window that contains a GLViewer instance.
class ViewerWindow(wx.Frame):
    def __init__(self, parent, **kwargs):
        title = "Multi-wavelength view"
        wx.Frame.__init__(self, parent, title = title, 
                style = wx.DEFAULT_FRAME_STYLE)

        sizer = wx.BoxSizer(wx.VERTICAL)
        self.viewer = GLViewer(self, size = (512, 512), **kwargs)

        sizer.Add(self.viewer)
        self.SetSizerAndFit(sizer)

        ## This lock prevents us from trying to update more than one
        # viewer at a time, evading some OpenGL errors.
        self.lock = threading.Lock()

        self.Show()


    ## Display a new image in the viewer, or hide the viewer if it's
    # been disconnected.
    def processNewImage(self, camId):
        if not self.viewer.haveInitedGL:
            return
        self.lock.acquire()

        data = numpy.random.random_integers(0, 10, (512, 512)).astype(numpy.uint16)

        self.viewer.updateImage(camId, data, 0, 10)

        color = colors[camId]
        # Remap from [0, 255] to [0, 1]
        color = [c / 255.0 for c in color]
        self.viewer.setColor(camId, color, False)

        self.viewer.Refresh()
        
        self.lock.release()
            


## OpenGL canvas for displaying multiple camera views layered on top of each
# other.
class GLViewer(wx.glcanvas.GLCanvas):
    ## Instantiate.
    def __init__(self, parent,
                 style = 0, size = wx.DefaultSize):
        wx.glcanvas.GLCanvas.__init__(self, parent, style = style, size = size)

        ## List of Image instances
        self.imgList = []
        for i in xrange(3):
            self.imgList.append(Image())

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

        wx.EVT_PAINT(self, self.OnPaint)
        # Do nothing on background erasure, to avoid flickering.
        wx.EVT_ERASE_BACKGROUND(self, lambda event: event)
    
        
    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


    def updateImage(self, index, data, imageMin, imageMax):
        self.pic_ny, self.pic_nx = data.shape
        self.imgList[index].updateImage(data, imageMin, imageMax)


    def setColor(self, imgidx, color, RefreshNow=1):
        self.imgList[imgidx].color = color
        if RefreshNow:
            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()

        glEnable(GL_TEXTURE_2D)
        glEnable(GL_BLEND)
        glBlendFunc(GL_ONE, GL_ONE)

        for image in self.imgList:
            if image.imageData is not None:
                image.render()

        glDisable(GL_TEXTURE_2D)
        glDisable(GL_BLEND)

        glFlush()
        glPopMatrix()
        self.SwapBuffers()


    def OnReload(self, event=None):
        self.Refresh(False)


app = wx.App(False)
window = ViewerWindow(None)

import time
def eventSpammer():
    while True:
        for i in xrange(3):
            event = wx.CommandEvent(wx.wxEVT_COMMAND_BUTTON_CLICKED, i)
            wx.PostEvent(window, event)
        time.sleep(.05)


def onNewImageReady(event):
    window.processNewImage(event.GetId())

# Listen to events
for i in xrange(3):
    wx.EVT_BUTTON(window, i, onNewImageReady)
print "Event listeners started"

threading.Thread(target = eventSpammer).start()
print "Spammer started"
app.MainLoop()