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")