import vtk, sys

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

class MainWindow(QtGui.QMainWindow):
    
    def __init__(self):
        super(MainWindow, self).__init__()
        
        self.createActions()
        self.createMenus()
        self.createQVTKRenderWidget()
        self.setWindowTitle('Basic Example To Use Python, Qt And VTK')
        self.resize(640, 480)
    
    def createActions(self):
        self.coneAct = QtGui.QAction('C&one', self, checkable=True, checked=True, triggered=self.coneSource)
        self.cubeAct = QtGui.QAction('C&ube', self, checkable=True, triggered=self.cubeSource)
        self.cylinderAct = QtGui.QAction('C&ylinder', self, checkable=True, triggered=self.cylinderSource)
        self.sphereAct = QtGui.QAction('S&phere', self, checkable=True, triggered=self.sphereSource)
        
        sourcesActGroup = QtGui.QActionGroup(self)
        sourcesActGroup.addAction(self.coneAct)
        sourcesActGroup.addAction(self.cubeAct)
        sourcesActGroup.addAction(self.cylinderAct)
        sourcesActGroup.addAction(self.sphereAct)
    
    def createMenus(self):
        sourcesMenu = self.menuBar().addMenu('&Sources')
        sourcesMenu.addAction(self.coneAct)
        sourcesMenu.addAction(self.cubeAct)
        sourcesMenu.addAction(self.cylinderAct)
        sourcesMenu.addAction(self.sphereAct)
    
    def coneSource(self):
        source = vtk.vtkConeSource()
        source.SetResolution(50)
        
        self.mapper.SetInput(source.GetOutput())
        self.actor.GetProperty().SetColor(1, 0, 0)
    
    def cubeSource(self):
        self.mapper.SetInput(vtk.vtkCubeSource().GetOutput())
        self.actor.GetProperty().SetColor(0, 1, 0)
    
    def cylinderSource(self):
        source = vtk.vtkCylinderSource()
        source.SetResolution(50)
        
        self.mapper.SetInput(source.GetOutput())
        self.actor.GetProperty().SetColor(0, 0, 1)
    
    def sphereSource(self):
        source = vtk.vtkSphereSource()
        source.SetPhiResolution(25)
        source.SetThetaResolution(50)
        
        self.mapper.SetInput(source.GetOutput())
        self.actor.GetProperty().SetColor(1, 1, 0)
    
    def createQVTKRenderWidget(self):
        self.mapper = vtk.vtkDataSetMapper()
        
        self.actor = vtk.vtkActor()
        self.actor.SetMapper(self.mapper)
        
        self.coneSource()
        
        renderer = vtk.vtkRenderer()
        renderer.AddActor(self.actor)
        
        interactor = QVTKRenderWindowInteractor(self)
        interactor.GetRenderWindow().AddRenderer(renderer)
        interactor.Initialize()
        interactor.Start()
        
        self.setCentralWidget(interactor)

if __name__ == '__main__':
    app = QtGui.QApplication(sys.argv)
    win = MainWindow()
    win.show()
    sys.exit(app.exec_())