content: """
  Example usage:
  $ python3 script/compress_data.py --dataset_dir /scr/lucyshi/dataset/aloha_test
  """
  import os
  import h5py
  import cv2
  import numpy as np
  import argparse
  from tqdm import tqdm
  
  # Constants
  DT = 0.02
  JOINT_NAMES = ["waist", "shoulder", "elbow", "forearm_roll", "wrist_angle", "wrist_rotate"]
  STATE_NAMES = JOINT_NAMES + ["gripper"]
  
  
  def compress_dataset(input_dataset_path, output_dataset_path):
      # Check if output path exists
      if os.path.exists(output_dataset_path):
          print(f"The file {output_dataset_path} already exists. Exiting...")
          return
  
      # Load the uncompressed dataset
      with h5py.File(input_dataset_path, 'r') as infile:
          # Create the compressed dataset
          with h5py.File(output_dataset_path, 'w') as outfile:
  
              outfile.attrs['sim'] = infile.attrs['sim']
              outfile.attrs['compress'] = True
  
              # Copy non-image data directly
              for key in infile.keys():
                  if key != 'observations':
                      outfile.copy(infile[key], key)
  
              obs_group = infile['observations']
  
              # Create observation group in the output
              out_obs_group = outfile.create_group('observations')
  
              # Copy non-image data in observations directly
              for key in obs_group.keys():
                  if key != 'images':
                      out_obs_group.copy(obs_group[key], key)
  
              image_group = obs_group['images']
              out_image_group = out_obs_group.create_group('images')
  
              # JPEG compression parameters
              encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 50]
  
              compressed_lens = []  # List to store compressed lengths for each camera
  
              for cam_name in image_group.keys():
                  if "_depth" in cam_name:  # Depth images are not compressed
                      out_image_group.copy(image_group[cam_name], cam_name)
                  else:
                      images = image_group[cam_name]
                      compressed_images = []
                      cam_compressed_lens = []  # List to store compressed lengths for this camera
  
                      # Compress each image
                      for image in images:
                          result, encoded_image = cv2.imencode('.jpg', image, encode_param)
                          compressed_images.append(encoded_image)
                          cam_compressed_lens.append(len(encoded_image))  # Store the length
  
                      compressed_lens.append(cam_compressed_lens)
  
                      # Find the maximum length of the compressed images
                      max_len = max(len(img) for img in compressed_images)
  
                      # Create dataset to store compressed images
                      compressed_dataset = out_image_group.create_dataset(cam_name, (len(compressed_images), max_len), dtype='uint8')
  
                      # Store compressed images
                      for i, img in enumerate(compressed_images):
                          compressed_dataset[i, :len(img)] = img
  
              # Save the compressed lengths to the HDF5 file
              compressed_lens = np.array(compressed_lens)
              _ = outfile.create_dataset('compress_len', compressed_lens.shape)
              outfile['/compress_len'][...] = compressed_lens
  
      print(f"Compressed dataset saved to {output_dataset_path}")
  
  
  def save_videos(video, dt, video_path=None):
      if isinstance(video, list):
          cam_names = list(video[0].keys())
          h, w, _ = video[0][cam_names[0]].shape
          w = w * len(cam_names)
          fps = int(1/dt)
          out = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
          # bitrate = 1000000
          # out.set(cv2.VIDEOWRITER_PROP_BITRATE, bitrate)
          for ts, image_dict in enumerate(video):
              images = []
              for cam_name in cam_names:
                  image = image_dict[cam_name]
                  image = image[:, :, [2, 1, 0]] # swap B and R channel
                  images.append(image)
              images = np.concatenate(images, axis=1)
              out.write(images)
          out.release()
          print(f'Saved video to: {video_path}')
      elif isinstance(video, dict):
          cam_names = list(video.keys())
          # Remove depth images
          cam_names = [cam_name for cam_name in cam_names if '_depth' not in cam_name]
          all_cam_videos = []
          for cam_name in cam_names:
              all_cam_videos.append(video[cam_name])
          all_cam_videos = np.concatenate(all_cam_videos, axis=2) # width dimension
  
          n_frames, h, w, _ = all_cam_videos.shape
          fps = int(1 / dt)
          out = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
          for t in range(n_frames):
              image = all_cam_videos[t]
              image = image[:, :, [2, 1, 0]]  # swap B and R channel
              out.write(image)
          out.release()
          print(f'Saved video to: {video_path}')
  
   
  def load_and_save_first_episode_video(dataset_dir, video_path):
      dataset_name = 'episode_0'
      _, _, _, _, image_dict = load_hdf5(dataset_dir, dataset_name)
      save_videos(image_dict, DT, video_path=video_path)
  
  
  def load_hdf5(dataset_dir, dataset_name):
      dataset_path = os.path.join(dataset_dir, dataset_name + '.hdf5')
      if not os.path.isfile(dataset_path):
          print(f'Dataset does not exist at \n{dataset_path}\n')
          exit()
  
      with h5py.File(dataset_path, 'r') as root:
          compressed = root.attrs.get('compress', False)
          image_dict = dict()
          for cam_name in root[f'/observations/images/'].keys():
              image_dict[cam_name] = root[f'/observations/images/{cam_name}'][()]
          if compressed:
              compress_len = root['/compress_len'][()]
  
      if compressed:
          for cam_id, cam_name in enumerate(image_dict.keys()):
              padded_compressed_image_list = image_dict[cam_name]
              image_list = []
              for frame_id, padded_compressed_image in enumerate(padded_compressed_image_list):
                  image_len = int(compress_len[cam_id, frame_id])
                  compressed_image = padded_compressed_image
                  image = cv2.imdecode(compressed_image, 1)
                  image_list.append(image)
              image_dict[cam_name] = image_list
  
      return None, None, None, None, image_dict  # Return only the image dict for this application
  
  
  if __name__ == '__main__':
      parser = argparse.ArgumentParser(description="Compress all HDF5 datasets in a directory.")
      parser.add_argument('--dataset_dir', action='store', type=str, required=True, help='Directory containing the uncompressed datasets.')
  
      args = parser.parse_args()
  
      output_dataset_dir = args.dataset_dir + '_compressed'
      os.makedirs(output_dataset_dir, exist_ok=True)
  
      # Iterate over each file in the directory
      for filename in tqdm(os.listdir(args.dataset_dir), desc="Compressing data"):
          if filename.endswith('.hdf5'):
              input_path = os.path.join(args.dataset_dir, filename)
              output_path = os.path.join(output_dataset_dir, filename)
              compress_dataset(input_path, output_path)
  
      # After processing all datasets, load and save the video for the first episode
      print(f'Saving video for episode 0 in {output_dataset_dir}')
      video_path = os.path.join(output_dataset_dir, 'episode_0_video.mp4')
      load_and_save_first_episode_video(output_dataset_dir, video_path)
  
  