import numpy as np
import vtk
# -------------------------
# Load data (x y z vx vy vz)
# -------------------------
data = np.loadtxt("./vector_field.txt")  # change path if needed
# If your file contains other columns (e.g. scalar) adjust indexing.
#
points = vtk.vtkPoints()
vectors = vtk.vtkFloatArray()
vectors.SetNumberOfComponents(3)
vectors.SetName("vectors")
#
magnitudes = vtk.vtkFloatArray()
magnitudes.SetNumberOfComponents(1)
magnitudes.SetName("Magnitude")
#
for row in data:
    x, y, z, vx, vy, vz = row
    points.InsertNextPoint(x, y, z)
    vectors.InsertNextTuple((vx, vy, vz))
    mag = (vx*vx + vy*vy + vz*vz) ** 0.5
    magnitudes.InsertNextValue(mag)
#
# Create polydata and attach vectors + magnitude scalars
polydata = vtk.vtkPolyData()
polydata.SetPoints(points)
polydata.GetPointData().SetVectors(vectors)
polydata.GetPointData().AddArray(magnitudes)        # add magnitude array
polydata.GetPointData().SetScalars(magnitudes)     # make it the active scalar array
#
# -------------------------
# Glyph (arrow) setup
# -------------------------
arrow = vtk.vtkArrowSource()
arrow.Update()
#
glyph = vtk.vtkGlyph3D()
glyph.SetSourceConnection(arrow.GetOutputPort())
glyph.SetInputData(polydata)
#
glyph.SetVectorModeToUseVector()      # orient by vector
glyph.OrientOn()
#
glyph.ScalingOn()
glyph.SetScaleModeToScaleByVector()  # scale by vector magnitude
glyph.SetScaleFactor(0.1)            # tweak to desired arrow length
glyph.SetColorModeToColorByScalar()  # color by scalar array (Magnitude)
glyph.Update()
#
# -------------------------
# Mapper & Actor
# -------------------------
mapper = vtk.vtkPolyDataMapper()
mapper.SetInputConnection(glyph.GetOutputPort())
#
# Tell mapper to use point field data and select the "Magnitude" array
mapper.SetScalarModeToUsePointFieldData()
mapper.SelectColorArray("Magnitude")
#
# Explicit scalar range from original magnitude array (safer)
mag_range = polydata.GetPointData().GetScalars("Magnitude").GetRange()
mapper.SetScalarRange(mag_range)
#
actor = vtk.vtkActor()
actor.SetMapper(mapper)
#
# -------------------------
# Scalar Bar
# -------------------------
scalar_bar = vtk.vtkScalarBarActor()
scalar_bar.SetLookupTable(mapper.GetLookupTable())
scalar_bar.SetTitle("Vector Magnitude")
scalar_bar.SetNumberOfLabels(5)
scalar_bar.SetMaximumWidthInPixels(150)
#
# -------------------------
# Renderer / Window / Interactor
# -------------------------
renderer = vtk.vtkRenderer()
renderer.AddActor(actor)
renderer.AddActor2D(scalar_bar)
renderer.SetBackground(0.1, 0.2, 0.4)
#
render_window = vtk.vtkRenderWindow()
render_window.AddRenderer(renderer)
render_window.SetSize(1200, 800)
#
interactor = vtk.vtkRenderWindowInteractor()
interactor.SetRenderWindow(render_window)
#
# Optional: nicer camera / lighting
renderer.ResetCamera()
render_window.Render()
interactor.Start()