import pandas as pd
from flowtracks.scene import Scene
import numpy as np

h5_name = r'../corner_51.h5'

particles = Scene(h5_name)
ids = particles.trajectory_ids()

dataframes = []
for id in particles.trajectory_ids()[:100]:
    traj  = particles.trajectory_by_id(id)

# or for traj in particles.iter_trajectories(): if you run for all trajectories
    dataframes.append(
        pd.DataFrame.from_records(
            traj, columns=["x", "y", "z", "dx", "dy", "dz", "frame", "particle"]
        )
    )

df = pd.concat(dataframes, ignore_index=True)
df["particle"] = df["particle"].astype(np.int32)
df["frame"] = df["frame"].astype(np.int32)
df.reset_index(inplace=True, drop=True)



# Paraview does not recognize it as a set without _000001.txt, so we the first 10000
# ptv_is.10001 is becoming ptv_00001.txt

print(df.head())
print('-------')
print(df.tail())

df_grouped = df.reset_index().groupby("frame")
for index, group in df_grouped:
    group.to_csv(
        f"./res/ptv_{int(index):05d}.txt",
        mode="w",
        columns=["particle", "x", "y", "z", "dx", "dy", "dz"],
        index=False,
    )

print("Saving trajectories to Paraview finished\n")


import pandas as pd
import vtk
from flowtracks.scene import Scene
import numpy as np

h5_name = r'../corner_51.h5'

particles = Scene(h5_name)
ids = particles.trajectory_ids()

dataframes = []
for id in particles.trajectory_ids()[:100]:
    traj = particles.trajectory_by_id(id)
    dataframes.append(
        pd.DataFrame.from_records(
            traj, columns=["x", "y", "z", "dx", "dy", "dz", "frame", "particle"]
        )
    )

df = pd.concat(dataframes, ignore_index=True)
df["particle"] = df["particle"].astype(np.int32)
df["frame"] = df["frame"].astype(np.int32)
df.reset_index(inplace=True, drop=True)

# Group by frame to create separate VTK files for animation
for frame_num, frame_data in df.groupby('frame'):
    # Create VTK points and point data
    points = vtk.vtkPoints()
    vectors = vtk.vtkFloatArray()
    vectors.SetNumberOfComponents(3)
    vectors.SetName("Velocity")
    
    particle_ids = vtk.vtkIntArray()
    particle_ids.SetName("ParticleID")
    
    # Add points and their associated data
    for _, row in frame_data.iterrows():
        points.InsertNextPoint(row['x'], row['y'], row['z'])
        vectors.InsertNextTuple3(row['dx'], row['dy'], row['dz'])
        particle_ids.InsertNextValue(row['particle'].astype(np.int32))
    
    # Create polydata
    polydata = vtk.vtkPolyData()
    polydata.SetPoints(points)
    polydata.GetPointData().AddArray(vectors)
    polydata.GetPointData().AddArray(particle_ids)
    
    # Write to VTK file
    writer = vtk.vtkXMLPolyDataWriter()
    filename = f"./vtk_output/particles_frame_{frame_num:05d}.vtp"
    writer.SetFileName(filename)
    writer.SetInputData(polydata)
    writer.Write()

print("VTK files have been created. Load them in Paraview as a time series.")