from PySide import QtGui, QtOpenGL, QtCore
from OpenGL import GL
from OpenGL import GL
from OpenGL.GL import shaders
from OpenGL.arrays import vbo
from OpenGL.GL.ARB import texture_rg
from OpenGL.GL.ARB import half_float_vertex
from ctypes import c_void_p
import numpy
import math

vert_src = """#version 120
attribute vec2 positionIn;
attribute vec2 texIn;
varying vec2 texcoord;
void main()
{
    gl_Position= vec4(positionIn * 2.0 - vec2(1),0,1);
    texcoord = texIn;
}
"""

frag_src = """#version 120
uniform sampler2D texture;
uniform int hdr_mode;
uniform float g;
uniform float m;
uniform float s;
varying vec2 texcoord;

float convert(float x)
{
    return clamp(pow(x*m,g) *s, 0.0, 1.0);
}

void main()
{
    vec2 coords = vec2(texcoord.x, 1.0 - texcoord.y);
    vec3 tex = texture2D(texture, coords).xyz;
    if (hdr_mode == 1) {
        gl_FragColor = vec4(convert(tex.x), convert(tex.y), convert(tex.z), 1);
    } else {
        gl_FragColor = vec4(tex,1);
    }
}
"""

class DisplayWidget(QtGui.QWidget):
    def __init__(self, parent):
        super(DisplayWidget, self).__init__(parent)
        
        self.buffers = None
        self.glWidget = GLWidget(self)       
              
        self.bufferComboBox = QtGui.QComboBox(self)
        policy = QtGui.QSizePolicy()
        policy.setHorizontalPolicy(QtGui.QSizePolicy.Expanding)
        self.bufferComboBox.setSizePolicy(policy)
        label = QtGui.QLabel("Buffers:")
        label.setBuddy(self.bufferComboBox)
        self.bufferComboBox.currentIndexChanged["QString"].connect(
            self.onBufferSelectChange)
        self.interactiveCheckBox = QtGui.QCheckBox("Interactive", self)
        midLayout = QtGui.QHBoxLayout()
        midLayout.addWidget(label)
        midLayout.addWidget(self.bufferComboBox)
        midLayout.addWidget(self.interactiveCheckBox)
                    
        self.label = QtGui.QLabel("Exposure: 0", self)
        self.slider = QtGui.QSlider(QtCore.Qt.Horizontal, self)
        self.slider.setRange(-100,100)
        self.slider.setValue(0)
        self.slider.valueChanged.connect(self.onExposureChange)
        bottomLayout = QtGui.QHBoxLayout()
        bottomLayout.addWidget(self.label)
        bottomLayout.addWidget(self.slider)
        bottomLayout.setSizeConstraint(QtGui.QLayout.SetMinimumSize)
        
        layout = QtGui.QVBoxLayout()
        layout.addWidget(self.glWidget)
        layout.addLayout(midLayout)
        layout.addLayout(bottomLayout)
        self.setLayout(layout)

    def setBuffers(self, buffers):
        for i in xrange(self.bufferComboBox.count()):
            self.bufferComboBox.removeItem(0)
        self.buffers = buffers
        buffersList = buffers.keys()
        buffersList.sort()
        self.bufferComboBox.addItems(buffersList)

    def setActiveBuffer(self, bufferName):
        idx = self.bufferComboBox.findText(bufferName)
        if idx == self.bufferComboBox.currentIndex():
            self.refreshDisplay()
        else:
            self.bufferComboBox.setCurrentIndex(idx)
   
    def onBufferSelectChange(self, value):
        if str(value) in self.buffers:
            ndarray = self.buffers[str(value)].data
            self.glWidget.copyDataToTexture(ndarray)
            if ndarray.dtype == numpy.float32 or ndarray.dtype == numpy.float16:
                self.slider.setEnabled(True)
            else:
                self.slider.setEnabled(False)
            self.glWidget.glDraw()

    def onExposureChange(self):
        value = 0.1 * self.slider.value()
        self.glWidget.exposure = value
        self.label.setText("Exposure: " + str(value))        
        self.glWidget.glDraw()

    def refreshDisplay(self):
        self.onBufferSelectChange(self.bufferComboBox.currentText())

    def sizeHint(self):
        return QtCore.QSize(400,400)

class GLWidget(QtOpenGL.QGLWidget):
    def __init__(self, parent):
        super(GLWidget, self).__init__(parent)
        
        self.w = 440
        self.h = 440

        self.rightBtnDown = False

        self.texture = None
        self.texturedata = None
        self.shader = None
        self.hdr_mode = 0
        self.vbo = None
        
        self.scale = 0.5
        self.steps = 0
        self.cx = 0.5
        self.cy = 0.5
        
        self.gamma = 1.0/2.2
        self.exposure = 0

        self.zoomFactor = 1.35
        self.panFactor = 0.002

    def initializeGL(self):
        pass

    def copyDataToTexture(self, ndarray):
        # Update dimensions of widget
        self.texturedata = ndarray
        self.w = ndarray.shape[0]
        self.h = ndarray.shape[1]
        self.updateGeometry()

        # Generate new texture
        if not self.texture:
            try:
                self.texture = GL.glGenTextures(1)
            except Exception:
                return
        target = GL.GL_TEXTURE_2D
        GL.glBindTexture(target, self.texture)
        GL.glTexParameterf(target, GL.GL_TEXTURE_MAG_FILTER, GL.GL_NEAREST)
        GL.glTexParameterf(target, GL.GL_TEXTURE_MIN_FILTER, GL.GL_NEAREST)
        GL.glTexParameterf(target, GL.GL_TEXTURE_WRAP_S, GL.GL_CLAMP_TO_EDGE)
        GL.glTexParameterf(target, GL.GL_TEXTURE_WRAP_T, GL.GL_CLAMP_TO_EDGE)
        GL.glTexParameteri(target, GL.GL_GENERATE_MIPMAP, GL.GL_FALSE);
        
        # Get texture format
        channels = ndarray.shape[2] if ndarray.ndim == 3 else 1
        if channels == 1:
            glFormat = GL.GL_RED
        elif channels == 2:
            glFormat = GL.GL_RG
        elif channels == 3:
            glFormat = GL.GL_RGB
        elif channels == 4:
            glFormat = GL.GL_RGBA
        glInternalFormat = glFormat

        # Get texture type
        if ndarray.dtype == numpy.float32:
            glType = GL.GL_FLOAT
            # Need to use the exposure shader if floating point
            self.hdr_mode = 1

            # The internal format changes with floating point textures
            if channels == 1:
                glInternalFormat = texture_rg.GL_R32F
            elif channels == 2:
                glInternalFormat = texture_rg.GL_RG32F
            elif channels == 3:
                glInternalFormat = GL.GL_RGB32F
            elif channels == 4:
                glInternalFormat = GL.GL_RGBA32F
        elif ndarray.dtype == numpy.float16:
            glType = GL.GL_FLOAT
            # Need to use the exposure shader if floating point
            self.hdr_mode = 1

            # The internal format changes with floating point textures
            if channels == 1:
                glInternalFormat = texture_rg.GL_R16F
            elif channels == 2:
                glInternalFormat = texture_rg.GL_RG16F
            elif channels == 3:
                glInternalFormat = GL.GL_RGB16F
            elif channels == 4:
                glInternalFormat = GL.GL_RGBA16F
        else:
            glType = GL.GL_UNSIGNED_BYTE
            self.hdr_mode = 0
        
        # Copy data to texture
        GL.glTexImage2D(target, 0, glInternalFormat, self.w, self.h,
                        0, glFormat, glType, ndarray)
        GL.glBindTexture(target, 0)
          
    def resizeGL(self, width, height):
        GL.glViewport(0,0,width,height)
        GL.glMatrixMode(GL.GL_PROJECTION)
        GL.glLoadIdentity()
        GL.glOrtho(0,1,0,1,0,1)
        GL.glMatrixMode(GL.GL_MODELVIEW)
      
    def compileShaders(self):
         # Build shaders
        vert_shader = shaders.compileShader(vert_src, GL.GL_VERTEX_SHADER)
        frag_shader = shaders.compileShader(frag_src, GL.GL_FRAGMENT_SHADER)
        self.shader = shaders.compileProgram(vert_shader, frag_shader)
        
    def paintGL(self):
        if GL.glCheckFramebufferStatus(GL.GL_FRAMEBUFFER) == 33305:
            return
        
        GL.glClear(GL.GL_COLOR_BUFFER_BIT)
        
        if not self.texture:
            return
            
        if not self.shader:
            self.compileShaders()

        if not self.vbo:
            self.vbo = GL.glGenBuffers(1)

        shaders.glUseProgram(self.shader)
               
        GL.glBindBuffer(GL.GL_ARRAY_BUFFER, self.vbo)