#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys
from vtk import *
from vtk.qt4.QVTKRenderWindowInteractor import QVTKRenderWindowInteractor
from PyQt4 import QtGui

class VTKWindow(QVTKRenderWindowInteractor):
    def __init__(self, parent, filename):
        super(VTKWindow, self).__init__(parent)

        self.parent = parent
        self.filename = filename

        self.reader = vtkUnstructuredGridReader()
        self.reader.SetFileName(self.filename)
        self.reader.Update()
        self.output = self.reader.GetOutput()
        self.scalar_range = self.output.GetScalarRange()
        
        # Create the custom lut and add it to the mapper
        self.lut = vtkLookupTable()
        self.lut.Build()
        
        self.mapper = vtkDataSetMapper()
        self.mapper.SetInput(self.output)
        self.mapper.SetScalarRange(self.scalar_range)
        self.mapper.SetLookupTable(self.lut)

        # Create the scalar_bar (legend) and set it up.
        self.scalar_bar = vtkScalarBarActor()
        self.scalar_bar.SetLookupTable(self.lut)
        
        self.actor = vtkActor()
        self.actor.SetMapper(self.mapper)

        self.renderer = vtkRenderer()
        self.renderer.AddActor(self.actor)
        self.renderer.AddActor2D(self.scalar_bar)

        self.renderer_window = self.GetRenderWindow()
        self.renderer_window.AddRenderer(self.renderer)
        self.renderer_window.SetSize(300,500)

        # In order to be able to move the scalar_bar we must use it as an
        # actor to vtkScalarBarWidget.
        # The vtkScalarBarWidget's 'On' method must be invoked last.
        self.scalar_bar_widget = vtkScalarBarWidget()
        self.scalar_bar_widget.SetScalarBarActor(self.scalar_bar)
        self.scalar_bar_widget.SetInteractor(self)
        self.scalar_bar_widget.On()
        
class VTKFrame(QtGui.QFrame):
    """Test class"""
    def __init__(self, parent, filename):
        super(VTKFrame, self).__init__(parent)

        hbox = QtGui.QHBoxLayout()

        # create the widget
        self.vtk_window = VTKWindow(self, filename)

        hbox.addWidget(self.vtk_window)
        self.setLayout(hbox)

if __name__ == "__main__":
     app = QtGui.QApplication(sys.argv)
     main = VTKFrame(None, "simple.vtk")
     main.show()
     sys.exit(app.exec_())