import vtk
from PySide6.QtWidgets import QApplication, QMainWindow, QVBoxLayout, QWidget
from vtkmodules.qt.QVTKRenderWindowInteractor import QVTKRenderWindowInteractor


class VolumeRenderingWindow(QMainWindow):
    def __init__(self):
        super().__init__()

        reader = vtk.vtkImageReader2()
        reader.SetDataExtent(0, 511, 0, 499, 0, 511)
        reader.SetDataByteOrderToLittleEndian()
        reader.SetDataScalarTypeToUnsignedShort()
        reader.SetFileName("D:/Downloads/christmas_tree_512x499x512_uint16.raw")  # Replace with your file path
        reader.Update()

        volume_mapper = vtk.vtkGPUVolumeRayCastMapper()  # You can also use vtkFixedPointVolumeRayCastMapper
        volume_mapper.SetInputData(reader.GetOutput())

        volume_property = vtk.vtkVolumeProperty()
        color_transfer_function = vtk.vtkColorTransferFunction()
        color_transfer_function.AddRGBPoint(0, 0, 0, 0)  # Black
        color_transfer_function.AddRGBPoint(255, 1, 1, 1)  # White
        volume_property.SetColor(color_transfer_function)

        gradient_opacity_transfer_function = vtk.vtkPiecewiseFunction()
        gradient_opacity_transfer_function.AddPoint(0, 0.0)
        gradient_opacity_transfer_function.AddPoint(90, 0.5)
        gradient_opacity_transfer_function.AddPoint(100, 1.0)
        volume_property.SetGradientOpacity(gradient_opacity_transfer_function)

        volume_actor = vtk.vtkVolume()
        volume_actor.SetMapper(volume_mapper)
        volume_actor.SetProperty(volume_property)

        renderer = vtk.vtkRenderer()
        renderer.AddVolume(volume_actor)

        # Create a PySide6 widget to embed the VTK rendering window
        vtk_widget = QVTKRenderWindowInteractor()
        vtk_widget.SetRenderWindow(renderer.GetRenderWindow())

        # Set up the main window layout
        central_widget = QWidget()
        layout = QVBoxLayout()
        layout.addWidget(vtk_widget)
        central_widget.setLayout(layout)
        self.setCentralWidget(central_widget)

        # Set background color
        renderer.SetBackground(0.1, 0.2, 0.4)


if __name__ == "__main__":
    app = QApplication([])
    window = VolumeRenderingWindow()
    window.show()
    app.exec()