import os
import sys
from glob import glob
import json
import argparse

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="")
    parser.add_argument("--env", default="AstroRoboSasa-Nes", help="env name")
    args = parser.parse_args()

    env = args.env

    json_in_path = "/home/zhiwen/projects/SymReL_original/data/jsons/{}.json".format(env)
    json_out_dir = "/home/zhiwen/projects/SymReL_original/data/{}-Json".format(env)
    os.makedirs(json_out_dir, exist_ok=True)

    with open(json_in_path, 'r') as f:
        data = json.load(f)

    for i, each_frame in enumerate(data):
        cnt = i + 1
        # cnt = i
        json_out_path = os.path.join(json_out_dir, f"frame_{cnt:04d}.json")
        # json_out_path = os.path.join(json_out_dir, f"frame_{cnt:05d}.json")
        with open(json_out_path, "w") as f:
            json.dump(each_frame, f, indent=4)