from OpenGL.GL import *
from OpenGL.arrays import vbo
from OpenGL.GL import shaders

import numpy
import wx
import wx.glcanvas


class MyCanvas(wx.glcanvas.GLCanvas):
    def __init__(self, parent):
        wx.glcanvas.GLCanvas.__init__(self, parent)
        wx.EVT_PAINT(self, self.OnPaint)
        wx.EVT_SIZE(self, lambda event: self.Refresh(0))
        self.haveInited = False
        self.image = None
        self.context = wx.glcanvas.GLContext(self)

    def OnPaint( self, mode):
        if not self.haveInited:
            self.SetCurrent(self.context)
            glClearColor(1, 1, 1, 1)

            self.displayTextureID = self.makeDisplayTexture()
            self.colorTextureID = self.makeColorTexture()

            vertex = shaders.compileShader("""
            void main() {
                gl_FrontColor = gl_Color;
                gl_TexCoord[0].xy = gl_MultiTexCoord0.xy;
                gl_Position = gl_ModelViewProjectionMatrix * gl_Vertex;
            }""", GL_VERTEX_SHADER)
            fragment = shaders.compileShader("""
            uniform sampler2D texture;
            uniform sampler1D color_lut;
            void main() {
                vec2 uv = gl_TexCoord[0].xy;
                vec4 color = texture2D(texture, uv);
                gl_FragColor = texture1D(color_lut, color.a);
            }""", GL_FRAGMENT_SHADER)
            self.shader = shaders.compileProgram(vertex, fragment)
            self.haveInited = True

        try:
            glClearColor(1, 1, 1, 1)
            glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
            glEnable(GL_BLEND)
            glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
            glColor(1, 1, 1, 1)

            glPushMatrix()
            glLoadIdentity()
            glTranslatef(-1, -1, 0)
            glScalef(1/256.0, 1/256.0, 1)

            glEnable(GL_TEXTURE_1D)
            glActiveTexture(GL_TEXTURE1)
            glBindTexture(GL_TEXTURE_1D, self.colorTextureID)

            glEnable(GL_TEXTURE_2D)
            glActiveTexture(GL_TEXTURE0)
            glBindTexture(GL_TEXTURE_2D, self.displayTextureID)

            shaders.glUseProgram(self.shader)
            loc = glGetUniformLocation(self.shader, 'texture')
            glUniform1i(loc, 0)
            loc = glGetUniformLocation(self.shader, 'color_lut')
            glUniform1i(loc, 1)

            glBegin(GL_QUADS)
            for x, y in [(0, 0), (1, 0), (1, 1), (0, 1)]:
                glVertex2f(x * 512, y * 512)
                glTexCoord2f(x, y)
            glEnd()

            shaders.glUseProgram(0)

            glPopMatrix()
            glFlush()
            self.SwapBuffers()
        except Exception, e:
            import traceback; traceback.print_exc()
            import sys; sys.exit()


    # Make the display texture, a uniform gradient.
    def makeDisplayTexture(self):
        data = numpy.zeros((512, 512), dtype = numpy.uint16)
        data[:] = numpy.linspace(0, 2 ** 15 - 1, 512)
        textureID = glGenTextures(1)
        glEnable(GL_TEXTURE_2D)
        glBindTexture(GL_TEXTURE_2D, textureID)
        glTexImage2D(GL_TEXTURE_2D, 0, GL_ALPHA32F_ARB, 512, 512, 0,
                        GL_ALPHA, GL_SHORT, data)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
        glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
        return textureID


    def makeColorTexture(self):
        # A band that goes from dark blue to cyan to yellow to red.
        colors = numpy.zeros((512, 3)).astype(numpy.float32)
        # Red channel
        colors[180:341, 0] = numpy.linspace(0, 1, 161)
        colors[341:441, 0] = 1
        colors[441:, 0] = numpy.linspace(1, .5, 71)
        # Green channel
        colors[51:200, 1] = numpy.linspace(0, 1, 149)
        colors[200:311, 1] = 1
        colors[311:461, 1] = numpy.linspace(1, 0, 150)
        # Blue channel
        colors[:51, 2] = numpy.linspace(.5, 1, 51)
        colors[51:160, 2] = 1
        colors[160:331, 2] = numpy.linspace(1, 0, 171)
        textureID = glGenTextures(1)
        glEnable(GL_TEXTURE_1D)
        glBindTexture(GL_TEXTURE_1D, textureID)
        glTexImage1D(GL_TEXTURE_1D, 0, GL_RGB, 512, 0, GL_RGB, GL_FLOAT,
                     colors)
        glTexParameterf(GL_TEXTURE_1D, GL_TEXTURE_MAG_FILTER, GL_NEAREST)
        glTexParameterf(GL_TEXTURE_1D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
        return textureID


class MyApp(wx.App):
    def OnInit(self):
        frame = wx.Frame(None)
        frame.SetSize((512, 512))
        canvas = MyCanvas(frame)
        frame.Show()
        return True
MyApp(redirect = False).MainLoop()