import time
import os
import numpy as np
import argparse
import matplotlib.pyplot as plt
import h5py
from roboimi.envs.vx300s_ee import make_ee_sim_env
from roboimi.envs.vx300s_jnt import make_sim_env
# from io_utils import IOUtils
from roboimi.utils.constants import PUPPET_GRIPPER_POSITION_NORMALIZE_FN, SIM_TASK_CONFIGS
# from ee_sim_env import make_ee_sim_env
from scripted_policy import PickAndTransferPolicy, InsertionPolicy,TestInsertionPolicy
import cv2

import pathlib
HOME_PATH = str(pathlib.Path(__file__).parent.resolve())
DATASET_DIR = HOME_PATH + '/dataset'


def main():
    task_name = 'sim_insertion'
    dataset_dir = DATASET_DIR + '/sim_insertion'   #SIM_TASK_CONFIGS[task_name]['dataset_dir']
    num_episodes = 20       #SIM_TASK_CONFIGS[task_name]['num_episodes']
    onscreen_render = None  #config['onscreen_render']
    inject_noise = False
    render_cam_name = 'angle'

    episode_len  = 1950      #SIM_TASK_CONFIGS[task_name]['episode_len']
    camera_names = ['top']   #SIM_TASK_CONFIGS[task_name]['camera_names']
    if task_name == 'sim_transfer_cube':
        print(task_name)
        # policy_cls = PickAndTransferPolicy
    elif task_name == 'sim_insertion':
        policy = TestInsertionPolicy(inject_noise)
        # policy_cls = InsertionPolicy
        print(task_name,'\n')
    else:
        raise NotImplementedError
    
    success = []
    
    
    for episode_idx in range(num_episodes):
        policy = TestInsertionPolicy(inject_noise)
        env_ee = make_ee_sim_env(task_name)
        print(f'\n{episode_idx=}')
        print('Rollout out EE space scripted policy')
        env_ee.reset()
        if onscreen_render:
            print("camera render")
        reward_ee = []
        jnt_traj = []
        gripper_ctrl_traj = []
        for step in range(episode_len):
            if step == 0:
                action = policy.predict(env_ee.obs,step)
                env_ee.step(action)
                env_ee.render()
            if step >= 1:
                action = policy.predict(env_ee.obs,step)
                env_ee.step(action)
                if step == 1:
                    box_pose = env_ee.obs['env_state'].copy()
                env_ee.render()
                reward_ee.append(env_ee.rew)
                jnt_traj.append(env_ee.obs['qpos'])
                gripper_ctrl_traj.append(env_ee.obs['gripper_ctrl'])
            # print("\nstep finish",step)
        sum_reward = np.sum(reward_ee)
        max_reward = np.max(reward_ee)
        if max_reward == env_ee.max_reward:
            print(f"{episode_idx=} Successful, {sum_reward=}")
        else:
            print(f"{episode_idx=} Failed")

        del policy
        # env_ee.viewer.close()
        del env_ee
        
        print('Replaying joint commands')
        time.sleep(1)
        env = make_sim_env(task_name)

        reward = []
        obs = []
        env.reset(box_pose)
        # print("envt=====",env.event)
        for step in range(2,len(jnt_traj)):
            action = jnt_traj[step]
            env.step(action)
            if step == 2:  #is np.ndarray
                print("step == 2 :type of env.obs.top===",type(env.obs['images']['top']))
            obs.append(env.obs)
            env.render()
            # print("\nstep===",step)

            reward.append(env.rew)
        sum_reward = np.sum(reward)
        max_reward = np.max(reward)
        # env.event.set()
        env.viewer.close()
        if max_reward == env.max_reward:
            success.append(1)

            print('Replaying',f"{episode_idx=} Successful, {sum_reward=}")
            t0 = time.time()
            data_dict = {
            '/observations/qpos': [],
            '/action': [],
            }
            
            for cam_name in camera_names:
                data_dict[f'/observations/images/{cam_name}'] = []
            for i in range(2,len(jnt_traj)):
                action = jnt_traj[i]
                data_dict['/observations/qpos'].append(obs[i-2]['qpos'])
                data_dict['/action'].append(action)
                for cam_name in camera_names:
                    data_dict[f'/observations/images/{cam_name}'].append(obs[i-2]['images'][cam_name])
            
            dataset_path = os.path.join(dataset_dir, f'episode_{episode_idx}')

            # with h5py.File(dataset_path + '.hdf5', 'w', rdcc_nbytes=1024 ** 2 * 2) as root:
            #     max_timesteps = len(jnt_traj)-2
            #     root.attrs['sim'] = True
            #     obs = root.create_group('observations')
            #     image = obs.create_group('images')
            #     for cam_name in camera_names:
            #         _ = image.create_dataset(cam_name, (max_timesteps, 480, 640, 3), dtype='uint8',
            #                              chunks=(1, 480, 640, 3), )
            #     qpos = obs.create_dataset('qpos', (max_timesteps, 14))
            #     action = root.create_dataset('action', (max_timesteps, 14))
                # 
                    # if name == '/observations/images/top':
                    #     root[name][...] = array
                    #     print("image done")
        #     # print(f'Saving: {time.time() - t0:.1f} secs\n')



        else:
            success.append(0)
            print('Replaying',f"{episode_idx=} Failed")



        
        del env

    print(f'Success: {np.sum(success)} / {len(success)}')
if __name__ == '__main__':
    main()