import sys

import vtk  # <-- enable this to fix the import error

from PySide6.QtWidgets import QApplication, QMainWindow, QFrame, QVBoxLayout

from vtkmodules.qt.QVTKRenderWindowInteractor import QVTKRenderWindowInteractor
from vtkmodules.vtkFiltersSources import vtkCubeSource
from vtkmodules.vtkRenderingCore import vtkRenderer, vtkActor, vtkPolyDataMapper


class CubeViewer(QMainWindow):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("VTK Cube Viewer")
        self.resize(800, 600)

        # Create central frame
        self.frame = QFrame()
        self.layout = QVBoxLayout()
        self.frame.setLayout(self.layout)
        self.setCentralWidget(self.frame)

        # Create VTK widget
        self.vtk_widget = QVTKRenderWindowInteractor(self.frame)
        self.layout.addWidget(self.vtk_widget)

        # Create renderer and render window
        self.renderer = vtkRenderer()
        self.renderer.SetBackground(0.2, 0.2, 0.3)  # Dark blue-gray background
        self.vtk_widget.GetRenderWindow().AddRenderer(self.renderer)

        # Initialize interactor
        self.interactor = self.vtk_widget.GetRenderWindow().GetInteractor()

        # Create cube
        self.create_cube()

        # Set camera position
        self.renderer.ResetCamera()
        camera = self.renderer.GetActiveCamera()
        camera.Elevation(30)
        camera.Azimuth(30)
        camera.Zoom(0.8)

        # Start interactor
        self.interactor.Initialize()

    def create_cube(self):
        # Create a cube source
        cube_source = vtkCubeSource()
        cube_source.SetXLength(1.0)
        cube_source.SetYLength(1.0)
        cube_source.SetZLength(1.0)
        cube_source.SetCenter(0.0, 0.0, 0.0)
        cube_source.Update()

        # Create a mapper
        mapper = vtkPolyDataMapper()
        mapper.SetInputConnection(cube_source.GetOutputPort())

        # Create an actor
        actor = vtkActor()
        actor.SetMapper(mapper)
        actor.GetProperty().SetColor(1.0, 0.6, 0.0)  # Orange color
        actor.GetProperty().SetAmbient(0.3)
        actor.GetProperty().SetDiffuse(0.7)
        actor.GetProperty().SetSpecular(0.5)
        actor.GetProperty().SetSpecularPower(20)

        # Add actor to renderer
        self.renderer.AddActor(actor)



    def closeEvent(self, event):
        # Properly clean up VTK interactor on window close
        self.vtk_widget.close()
        super().closeEvent(event)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = CubeViewer()
    window.show()
    sys.exit(app.exec())