import os,cv2
import numpy as np
import h5py
from roboimi.envs.double_pos_ctrl_env import make_sim_env

pwd = os.getcwd()
DATASET_DIR = pwd + '/demos/dataset'

def traverse_datasets(hdf_file):
    all_actions = []
    all_images = []
    def h5py_dataset_iterator(g, prefix=''):
        for key in g.keys():
            item = g[key]
            path = '{}/{}'.format(prefix, key)
            if isinstance(item, h5py.Dataset):  
                yield (path, item)
            elif isinstance(item, h5py.Group):  
                yield from h5py_dataset_iterator(item, path)

    with h5py.File(hdf_file, 'r') as f:
        i = 0
        for (path, dset) in h5py_dataset_iterator(f):
            i+=1
            if path == "/action":
                # for i in range():
                action = dset[()] 
                # action = dset 
                print("len action=",len(action[0]))
                print("action=",action[0])
                # all_actions.append(action)
            if path == "/observations/images/top":
                images = dset[()]  
                # all_images.append(images)
        print("i=",i)
    return all_actions,all_images 
if __name__ == "__main__":
    # env = make_sim_env('sim_transfer')
    dataset_dir = DATASET_DIR + '/sim_transfer/episode_0.hdf5'
    traverse_datasets(dataset_dir)
    # actions,images = traverse_datasets(dataset_dir)
    # print("len actions=",len(actions),"len image=",len(images))

    # box_pos = np.array([0.4878127780735517, 0.7978437132112287, 0.47])
    # env.mj_data.joint('red_box_joint').qpos[0] = box_pos[0] 
    # env.mj_data.joint('red_box_joint').qpos[1] = box_pos[1]
    # env.mj_data.joint('red_box_joint').qpos[2] = box_pos[2]
    # env.mj_data.joint('red_box_joint').qpos[3] = 1.0
    # env.mj_data.joint('red_box_joint').qpos[4] = 0.0
    # env.mj_data.joint('red_box_joint').qpos[5] = 0.0
    # env.mj_data.joint('red_box_joint').qpos[6] = 0.0
    # env.reset()
    # cv2.namedWindow('Cam view',cv2.WINDOW_NORMAL)
    
    # for i in range(len(actions)):
    #     action_l = actions[i][0:7]
    #     action_r = actions[i][7:14]
    #     gripper = np.array([actions[i][-2],actions[i][-1]])
    #     action = np.hstack((action_l,action_r,gripper))
    #     env.step(action)
    #     cv2.imshow('Cam view', images[i])
    #     cv2.waitKey(1)
    #     if env.is_render:
    #         env.render()