import os
import pickle


def main(file: str, trajectory_index=0, save_path="mesh", noise_level: float = 0.0):
    with open(file, 'rb') as f:
        data = pickle.load(f)
    # data is a list of dictionaries, each containing one trajectory

    trajectory = data[trajectory_index]
    save_path = os.path.join(save_path, f"trajectory_{trajectory_index:03d}", f"noise_{noise_level}")
    to_obj(trajectory, save_path=save_path, noise_level=noise_level)


def to_obj(trajectory: dict, save_path: str = "mesh", noise_level: float = 0.0):
    os.makedirs(save_path, exist_ok=True)

    vertices_over_time = trajectory["nodes_grid"]
    faces = trajectory["triangles_grid"]

    # add cumulate noise on the vertices over the trajectory
    # draw noise from a Gaussian with mean 0 and std noise_level
    import numpy as np
    noise = np.random.normal(0, noise_level, vertices_over_time.shape)
    # add **cumulate** noise to the vertices
    cumulate_noise = np.cumsum(noise, axis=0)
    vertices_over_time += cumulate_noise

    for time_step, vertices in enumerate(vertices_over_time):
        filename = os.path.join(save_path, f"mesh_{time_step:03d}.obj")
        with open(filename, 'w') as file:
            for v in vertices:
                file.write(f"v {v[0]} {v[1]} {v[2]}\n")
            for f in faces:
                # OBJ files are 1-indexed
                file.write(f"f {f[0] + 1} {f[1] + 1} {f[2] + 1}\n")


if __name__ == '__main__':
    main(file="../../datasets/lts_gns/teddy_fall_nopc/teddy_fall_nopc_test.pkl",
         trajectory_index=0,
         save_path="objs",
         noise_level=0)
