import vtk
from mayavi import mlab
import numpy as np

def create_disk(x, y, z, radius):
    # Create a disk (circle) in the xy-plane
    disk_source = vtk.vtkDiskSource()
    disk_source.SetInnerRadius(0)
    disk_source.SetOuterRadius(radius)
    disk_source.SetCircumferentialResolution(30)
    disk_source.SetRadialResolution(1)
    disk_source.Update()
    
    # Move the disk to the specified position in 3D space
    transform = vtk.vtkTransform()
    transform.Translate(x, y, z)
    
    transform_filter = vtk.vtkTransformPolyDataFilter()
    transform_filter.SetInputData(disk_source.GetOutput())
    transform_filter.SetTransform(transform)
    transform_filter.Update()
    
    return transform_filter.GetOutput()

def create_tube_from_disks(disks):
    # Combine multiple disks to form the tube mesh
    points = vtk.vtkPoints()
    lines = vtk.vtkCellArray()
    
    point_id = 0
    for disk in disks:
        num_points = disk.GetNumberOfPoints()
        for i in range(num_points):
            point = disk.GetPoint(i)
            points.InsertNextPoint(point)
        
        line = vtk.vtkLine()
        line.GetPointIds().SetId(0, point_id)
        line.GetPointIds().SetId(1, point_id + num_points - 1)
        lines.InsertNextCell(line)
        
        point_id += num_points
    
    tube_mesh = vtk.vtkPolyData()
    tube_mesh.SetPoints(points)
    tube_mesh.SetLines(lines)
    
    return tube_mesh

# Create three disks at different positions
disk1 = create_disk(0, 0, 0, 0.1)
disk2 = create_disk(0, 0, 1, 0.1)
disk3 = create_disk(0, 0, 2, 0.1)

# Combine the disks and create the tube mesh
disks = [disk1, disk2, disk3]
tube_mesh = create_tube_from_disks(disks)
# Extract vertices from the tube mesh
vertices = np.array(tube_mesh.GetPoints().GetData())

# Extract lines from the tube mesh
cell_array = tube_mesh.GetLines()
cell_array.InitTraversal()
ids = vtk.vtkIdList()
lines = []
while cell_array.GetNextCell(ids):
    line = [ids.GetId(i) for i in range(ids.GetNumberOfIds())]
    lines.append(line)

# Plot the tube mesh using Mayavi
mlab.figure(bgcolor=(1, 1, 1))
for line in lines:
    x = vertices[line, 0]
    y = vertices[line, 1]
    z = vertices[line, 2]
    mlab.plot3d(x, y, z, tube_radius=0.05, color=(0, 0, 1))

mlab.show()