if __name__ == "__main__":
    import sys
    import os
    import pathlib

    ROOT_DIR = str(pathlib.Path(__file__).parent.parent.parent)
    sys.path.append(ROOT_DIR)


import os
import click
import pathlib
import numpy as np
from diffusion_policy.common.replay_buffer import ReplayBuffer

@click.command()
@click.option('-i', '--input', required=True, help='input dir contains npy files')
@click.option('-o', '--output', required=True, help='output zarr path')
@click.option('--abs_action', is_flag=True, default=False)
def main(input, output, abs_action):
    data_directory = pathlib.Path(input)
    observations = np.load(
        data_directory / "multimodal_push_observations.npy"
    )
    actions = np.load(data_directory / "multimodal_push_actions.npy")
    masks = np.load(data_directory / "multimodal_push_masks.npy")

    buffer = ReplayBuffer.create_empty_numpy()
    for i in range(len(masks)):
        eps_len = int(masks[i].sum())
        obs = observations[i,:eps_len].astype(np.float32)
        action = actions[i,:eps_len].astype(np.float32)
        if abs_action:
            prev_eef_target = obs[:,8:10]
            next_eef_target = prev_eef_target + action
            action = next_eef_target
        data = {                              
            'obs': obs,
            'action': action
        }
        buffer.add_episode(data)

    buffer.save_to_path(zarr_path=output, chunk_length=-1)

if __name__ == '__main__':
    main()