# ParaView 6.x: Motive CSV -> time-varying marker animation
# Put this file in the SAME folder as the Motive CSV and rename the CSV to:
#     mocap.csv
# Then in ParaView: View > Python Shell > Run Script

from paraview.simple import *
import os

# ===== Settings =====
AMPLIFICATION = 50.0   # 1.0 = actual motion, 50.0 = displacement shown 50x larger
TUBE_RADIUS_MM = 8.0   # visual thickness of the line connecting markers
POINT_SIZE = 10.0
# ====================

try:
    SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
except NameError:
    SCRIPT_DIR = os.getcwd()

CSV_PATH = os.path.join(SCRIPT_DIR, "mocap.csv")

if not os.path.exists(CSV_PATH):
    raise RuntimeError(
        "mocap.csv was not found.\n"
        "Put this script in the same folder as the Motive CSV, "
        "and rename the CSV to mocap.csv."
    )

request_info = f"""
import csv

filename = {CSV_PATH!r}
times = []

with open(filename, "r", encoding="utf-8-sig", newline="") as f:
    reader = csv.reader(f)

    # Motive CSV:
    # 1 metadata line + 1 blank line + 6 header lines = 8 lines before data
    for _ in range(8):
        next(reader, None)

    for row in reader:
        if len(row) < 2:
            continue
        try:
            times.append(float(row[1]))
        except (ValueError, TypeError):
            pass

executive = self.GetExecutive()
outInfo = executive.GetOutputInformation(0)

outInfo.Remove(executive.TIME_STEPS())
for t in times:
    outInfo.Append(executive.TIME_STEPS(), t)

outInfo.Remove(executive.TIME_RANGE())
if times:
    outInfo.Append(executive.TIME_RANGE(), times[0])
    outInfo.Append(executive.TIME_RANGE(), times[-1])
"""

main_script = f"""
import csv
import bisect
import vtk

filename = {CSV_PATH!r}
amplification = {AMPLIFICATION!r}

# Cache the whole CSV after the first execution.
try:
    _mocap_cache
except NameError:
    _mocap_cache = {{}}

if filename not in _mocap_cache:
    times = []
    frames = []

    with open(filename, "r", encoding="utf-8-sig", newline="") as f:
        reader = csv.reader(f)
        for _ in range(8):
            next(reader, None)

        for row in reader:
            if len(row) < 32:
                continue

            try:
                t = float(row[1])
                coords = [float(v) for v in row[2:32]]
            except (ValueError, TypeError):
                continue

            times.append(t)
            frames.append(coords)

    if not frames:
        raise RuntimeError("No Motive marker data could be read from the CSV.")

    # 10 markers, each with X,Y,Z.
    n_markers = len(frames[0]) // 3

    # Connect markers in their initial X-coordinate order.
    order = sorted(range(n_markers), key=lambda j: frames[0][3*j])

    _mocap_cache[filename] = (times, frames, order)

times, frames, order = _mocap_cache[filename]

executive = self.GetExecutive()
outInfo = executive.GetOutputInformation(0)

if outInfo.Has(executive.UPDATE_TIME_STEP()):
    req_time = outInfo.Get(executive.UPDATE_TIME_STEP())
else:
    req_time = times[0]

# Choose the nearest measured time step.
idx = bisect.bisect_left(times, req_time)
if idx >= len(times):
    idx = len(times) - 1
elif idx > 0 and abs(times[idx-1] - req_time) <= abs(times[idx] - req_time):
    idx -= 1

current = frames[idx]
initial = frames[0]

points = vtk.vtkPoints()
disp = vtk.vtkDoubleArray()
disp.SetName("Displacement_mm")
disp.SetNumberOfComponents(3)

marker_index = vtk.vtkIntArray()
marker_index.SetName("MarkerIndex")

for out_i, j in enumerate(order):
    x0 = initial[3*j]
    y0 = initial[3*j + 1]
    z0 = initial[3*j + 2]

    x = current[3*j]
    y = current[3*j + 1]
    z = current[3*j + 2]

    dx = x - x0
    dy = y - y0
    dz = z - z0

    # Keep the initial geometry, and amplify only the measured displacement.
    xa = x0 + amplification * dx
    ya = y0 + amplification * dy
    za = z0 + amplification * dz

    points.InsertNextPoint(xa, ya, za)
    disp.InsertNextTuple3(dx, dy, dz)
    marker_index.InsertNextValue(j)

polyline = vtk.vtkPolyLine()
polyline.GetPointIds().SetNumberOfIds(len(order))
for i in range(len(order)):
    polyline.GetPointIds().SetId(i, i)

lines = vtk.vtkCellArray()
lines.InsertNextCell(polyline)

output = self.GetPolyDataOutput()
output.SetPoints(points)
output.SetLines(lines)
output.GetPointData().AddArray(disp)
output.GetPointData().AddArray(marker_index)
output.GetInformation().Set(output.DATA_TIME_STEP(), times[idx])
"""

# Remove previously-created object with the same name if this script is rerun.
for name in ("MocapMarkers", "MocapTube"):
    old = FindSource(name)
    if old is not None:
        Delete(old)

src = ProgrammableSource(registrationName="MocapMarkers")
src.OutputDataSetType = "vtkPolyData"
src.ScriptRequestInformation = request_info
src.Script = main_script
src.UpdatePipelineInformation()
src.UpdatePipeline()

# Marker display
marker_display = Show(src)
marker_display.Representation = "Points"
marker_display.PointSize = POINT_SIZE

# A tube makes the connected marker line easier to see.
tube = Tube(registrationName="MocapTube", Input=src)
tube.Radius = TUBE_RADIUS_MM
tube.NumberofSides = 12
tube.UpdatePipeline()

tube_display = Show(tube)
tube_display.Representation = "Surface"

# Animation setup
scene = GetAnimationScene()
scene.UpdateAnimationUsingDataTimeSteps()
scene.PlayMode = "Snap To TimeSteps"

ResetCamera()
Render()

print("Mocap animation loaded.")
print("CSV:", CSV_PATH)
print("Amplification:", AMPLIFICATION)
print("Use the Play button in the animation toolbar.")