import sys
sys.path.append('./')
import os
import sapien.core as sapien
from collections import OrderedDict
import pdb
from envs import *
import yaml
import importlib

def class_decorator(task_name):
    envs_module = importlib.import_module(f'envs.{task_name}')
    try:
        env_class = getattr(envs_module, task_name)
        env_instance = env_class()
    except:
        raise SystemExit("No such task")
    return env_instance

def main():
    task_name = input()
    
    task = class_decorator(task_name)
    task_config_path = f'./task_config/{task_name}.yml'
    if not os.path.isfile(task_config_path):
    # if True:
        data = {
            'task_name': task_name,
            'render_freq': 0,
            'use_seed': False,
            'collect_data': True,
            'save_path': './data',
            'dual_arm': True,
            'st_episode': 0 ,
            'camera_w': 320,
            'camera_h': 240,
            'pcd_crop': True ,
            'pcd_down_sample_num': 1024,
            'pose_type': "gt",
            'episode_num': 100,
            'save_type':{
                'raw_data': False,
                'pkl': True
            },
            'data_type':{
                'rgb': False,
                'observer': False,
                'depth': False,
                'pointcloud': True,
                'conbine': False,
                'endpose': True,
                'qpos': True,
                'mesh_segmentation': False,
                'actor_segmentation': False,
            }
        }
        with open(task_config_path, 'w') as f:
            yaml.dump(data,f,default_flow_style = False,sort_keys=False)
    with open(task_config_path, 'r', encoding='utf-8') as f:
        args = yaml.load(f.read(), Loader=yaml.FullLoader)

    run(task, args)


def run(Demo_class, args):
    epid = 0       
    seed_list=[]   
    suc_num = 0    
    fail_num = 0   

    if not args['use_seed']:
        while suc_num < args['episode_num']:
            try:
                Demo_class.setup_demo(now_ep_num=suc_num, seed = epid, **args)
                Demo_class.play_once()

                if Demo_class.plan_success and Demo_class.check_success():
                    print(f"simulate data episode {suc_num} success! (seed = {epid})")
                    seed_list.append(epid)
                    suc_num+=1
                else:
                    print(f"simulate data episode {suc_num} fail! (seed = {epid})   ")
                    fail_num +=1
                
                Demo_class.close()
                if (args['render_freq']):
                    Demo_class.viewer.close()
                epid +=1
            except Exception as e:
                import traceback
                traceback.print_exc()
                print(e)
                print(f"simulate data episode {suc_num} fail! (seed = {epid})   ")
                fail_num +=1
                Demo_class.close()
                if (args['render_freq']):
                    Demo_class.viewer.close()
                epid +=1

        os.makedirs('./task_config/seeds/', exist_ok=True)
        with open('./task_config/seeds/'+args['task_name']+'.txt', 'w') as file:
            for sed in seed_list:
                file.write("%s " % sed)
        print(f'\nComplete simulation, failed {fail_num} times')

    else:
        print(f'using saved seeds list')
        with open('./task_config/seeds/'+args['task_name']+'.txt', 'r') as file:
            seed_list = file.read().split()
            seed_list = [int(i) for i in seed_list]

    if args['collect_data']:
        print('Start data collection')

        args['render_freq']=0
        args['is_save'] = True
        for id in range(args['st_episode'], args['episode_num']):
            Demo_class.setup_demo(now_ep_num=id, seed = seed_list[id],**args)
            Demo_class.play_once()
            if Demo_class.save_type.get('raw_data', True):
                head_config = Demo_class.get_camera_config(Demo_class.head_camera)
                left_config = Demo_class.get_camera_config(Demo_class.left_camera)
                right_config = Demo_class.get_camera_config(Demo_class.right_camera)
                save_json(Demo_class.file_path["f_color"]+"config.json", head_config)
                save_json(Demo_class.file_path["l_color"]+"config.json", left_config)
                save_json(Demo_class.file_path["r_color"]+"config.json", right_config)

                save_json(Demo_class.file_path["f_depth"]+"config.json", head_config)
                save_json(Demo_class.file_path["l_depth"]+"config.json", left_config)
                save_json(Demo_class.file_path["r_depth"]+"config.json", right_config)

                save_json(Demo_class.file_path["f_pcd"]+"config.json", head_config)
                save_json(Demo_class.file_path["l_pcd"]+"config.json", left_config)
                save_json(Demo_class.file_path["r_pcd"]+"config.json", right_config)

            Demo_class.close()
            print('\nsuccess!')

            
if __name__ == "__main__":
    main()