## 本文件用于将robotwin Challenge 2 中的hdf5数据转为TinyVLA可以直接训练的数据。
import sys

sys.path.append('./policy/ACT/')

import os
import h5py
import numpy as np
import cv2
import argparse
import json

task_prompt = {
"place_object_scale": "Place the object onto the scale.",
"place_phone_stand": "Place phone onto stand using multi-angle desk images to determine positions and plan actions.",
"stack_blocks_three": "Move the red blocks to center of the table and stack the green block on the red block ,and stack the blue block on the green block.",
}
task_reasoning = {
    "place_object_scale": 0,
    "stack_blocks_three": 1
}
all_reasoning = [
    ["Pick up the object and place the object onto the scale."],
    ["Move the red blocks to center of the table and stack the green block on the red block ,and stack the blue block on the green block."],
]

def load_hdf5(dataset_path):
    '''
    从robotwin Challenge 2 生成的 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:
        left_gripper, left_arm = root['/joint_action/left_gripper'][()], root['/joint_action/left_arm'][()]
        right_gripper, right_arm = root['/joint_action/right_gripper'][()], root['/joint_action/right_arm'][()]
        image_dict = dict()  
        for cam_name in root[f'/observation/'].keys():
            image_dict[cam_name] = root[f'/observation/{cam_name}/rgb'][()]

    return left_gripper, left_arm, right_gripper, right_arm, image_dict



def data_transform(path, episode_num, save_path, task_name):
    '''
    将原始数据转换为 VLA 模型可以使用的格式，并保存为新的 HDF5 文件。
    '''
    begin = 0
    floders = os.listdir(path)  
    assert episode_num <= len(floders), "data num not enough"

    if not os.path.exists(save_path):
        os.makedirs(save_path)

    for i in range(episode_num):
        left_gripper_all, left_arm_all, right_gripper_all, right_arm_all, image_dict = load_hdf5(
            os.path.join(path, f"episode{i}.hdf5"))
        qpos = []
        actions = []
        cam_high = []
        cam_right_wrist = []
        cam_left_wrist = []
        left_arm_dim = []
        right_arm_dim = []

        last_state = None
        len_traj = left_gripper_all.shape[0]-1 # reasonging action obs的长度
        for j in range(0, left_gripper_all.shape[0]):

            left_gripper, left_arm, right_gripper, right_arm = left_gripper_all[j], left_arm_all[j], right_gripper_all[
                j], right_arm_all[j],

            if j != left_gripper_all.shape[0] - 1:
                state = np.concatenate((left_arm, [left_gripper], right_arm, [right_gripper]), axis=0)  # joint

                state = state.astype(np.float32)
                qpos.append(state)

                camera_high_bits = image_dict['head_camera'][j]
                camera_high = cv2.imdecode(np.frombuffer(camera_high_bits, np.uint8), cv2.IMREAD_COLOR)
                cam_high.append(camera_high)

                camera_right_wrist_bits = image_dict['right_camera'][j]
                camera_right_wrist = cv2.imdecode(np.frombuffer(camera_right_wrist_bits, np.uint8), cv2.IMREAD_COLOR)
                cam_right_wrist.append(camera_right_wrist)

                camera_left_wrist_bits = image_dict['left_camera'][j]
                camera_left_wrist = cv2.imdecode(np.frombuffer(camera_left_wrist_bits, np.uint8), cv2.IMREAD_COLOR)
                cam_left_wrist.append(camera_left_wrist)

            if j != 0:
                action = state
                actions.append(action)
                left_arm_dim.append(left_arm.shape[0])
                right_arm_dim.append(right_arm.shape[0])

        hdf5path = os.path.join(save_path, f'episode_{i}.hdf5')

        with h5py.File(hdf5path, 'w') as f:
            f.create_dataset('action', data=np.array(actions))
            language_raw = task_prompt[task_name].encode('utf-8')
            sub_reasons = [all_reasoning[task_reasoning[task_name]][0]] * int(len_traj)
            f.create_dataset('language_raw', data=np.array(language_raw)) # 增加指令
            f.create_dataset('reasoning', data=np.array(sub_reasons, dtype=object)) # 加载设定的推理
            obs = f.create_group('observations')
            obs.create_dataset('qpos', data=np.array(qpos))
            obs.create_dataset('qvel', data=np.array(qpos)) # 无意义为了对齐key
            obs.create_dataset('left_arm_dim', data=np.array(left_arm_dim))
            obs.create_dataset('right_arm_dim', data=np.array(right_arm_dim))
            image = obs.create_group('images')
            image.create_dataset('cam_high', data=np.stack(cam_high), dtype=np.uint8)
            image.create_dataset('cam_right_wrist', data=np.stack(cam_right_wrist), dtype=np.uint8)
            image.create_dataset('cam_left_wrist', data=np.stack(cam_left_wrist), dtype=np.uint8)

        begin += 1
        print(f"proccess {i} success!")

    return begin


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Process some episodes.')
    parser.add_argument('task_name', type=str, default='bottle_adjust',
                        help='The name of the task (e.g., bottle_adjust)')
    parser.add_argument('setting', type=str)
    parser.add_argument('expert_data_num', type=int, default=50,
                        help='Number of episodes to process (e.g., 50)')

    args = parser.parse_args()

    task_name = args.task_name
    setting = args.setting
    expert_data_num = args.expert_data_num

    data_path_name = task_name + "/" + setting + "/data"
    begin = 0
    begin = data_transform(os.path.join("../../data/", data_path_name), expert_data_num,
                           f"data/sim-{task_name}/{setting}-{expert_data_num}",task_name)
