import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

os.environ["MUJOCO_GL"] = "egl"
#os.environ["EGL_VISIBLE_DEVICES"] = "6"

from termcolor import colored
import subprocess

DATA_ROOT = './data/metaworld'

TASK_AGENT_CONFIG = [
    ('dial-turn', 'lift3d_bnmlp', 16),
]

CAMERAS = [
    'corner',
]


def test_codebase():
    for task, agent, batch_size in TASK_AGENT_CONFIG:
        for camera in CAMERAS:
            cmd = [
                'python', '-m', 'lift3d.tools.train_policy',
                '--config-name=train_metaworld',
                f'agent={agent}',
                f'task_name={task}',
                f'camera_name={camera}',
                f'dataloader.batch_size={batch_size}',
                f'dataset_dir={DATA_ROOT}/{task}_{camera}.zarr',
            ]
            print(colored('[INFO]', 'blue'), ' '.join(cmd))
            subprocess.run(cmd)



if __name__ == '__main__':
    test_codebase()
