from gym.envs.registration import register
import gym

MODEL_VERSION = 'VAR++' # choose from E2E, VAR, VAR++
if MODEL_VERSION == 'E2E':
    model_number = 1
elif MODEL_VERSION == 'VAR':
    model_number = 2
elif MODEL_VERSION == 'VAR++':
    model_number = 3
else:
    raise NotImplementedError

ENV='ai2thor' # choose from ai2thor, kuka, turtlebot, kinova_gen3, ROS

def main_config():
    config=None

    if model_number == 1:
        if ENV == 'ai2thor':
            from Envs.ai2thor.RSI1.config import AI2ThorConfig
            config = AI2ThorConfig()

        elif ENV == 'kuka':
            from Envs.pybullet.kuka.RSI1.config import KukaConfig
            config = KukaConfig()

        elif ENV == 'turtlebot':
            from Envs.pybullet.turtlebot.RSI1.config import TurtleBotConfig
            config = TurtleBotConfig()

        else:
            raise NotImplementedError

    elif model_number == 2:
        if ENV == 'ai2thor':
            from Envs.ai2thor.RSI2.config import AI2ThorConfig
            config = AI2ThorConfig()

        elif ENV == 'kuka':
            from Envs.pybullet.kuka.RSI2.config import KukaConfig
            config = KukaConfig()

        elif ENV == 'turtlebot':
            from Envs.pybullet.turtlebot.RSI2.config import TurtleBotConfig
            config = TurtleBotConfig()

        elif ENV == 'kinova_gen3':
            from Envs.pybullet.kinova_gen3.RSI2.config import KinovaGen3Config
            config = KinovaGen3Config()

        elif ENV=='ROS':
            from Envs.pybullet.kinova_gen3.ROS.config import KinovaGen3Config
            config = KinovaGen3Config()
        
        else:
            raise NotImplementedError

    elif model_number == 3:
        if ENV == 'ai2thor':
            from Envs.ai2thor.RSI3.config import AI2ThorConfig
            config = AI2ThorConfig()

        elif ENV == 'kuka':
            from Envs.pybullet.kuka.RSI3.config import KukaConfig
            config = KukaConfig()

        elif ENV == 'turtlebot':
            from Envs.pybullet.turtlebot.RSI3.config import TurtleBotConfig
            config = TurtleBotConfig()

        elif ENV == 'kinova_gen3':
            from Envs.pybullet.kinova_gen3.RSI3.config import KinovaGen3Config
            config = KinovaGen3Config()

        elif ENV == 'ROS':
            from Envs.pybullet.kinova_gen3.ROS.config import KinovaGen3Config
            config = KinovaGen3Config()

        else:
            raise NotImplementedError

    return config

def gym_register(config):

    env_folder = '' if ENV == 'ai2thor' else '.pybullet'

    entry_point = 'Envs' + env_folder + '.' + ENV + '.RSI' + str(model_number) + '.pretext_env_RSI' + str(
        model_number) + ':PretextEnvRSI' + str(model_number)
    register(
        id= ENV + '-pretext-v' + str(model_number),  # e.g. 'turtlebot-pretext-v1',
        entry_point=entry_point
    )
    entry_point = 'Envs' + env_folder + '.' + ENV + '.RSI' + str(model_number) + '.RL_env_RSI' + str(
        model_number) + ':RLEnvRSI' + str(model_number)
    register(
        id=ENV + '-RL-v' + str(model_number),  # e.g. turtlebot-RL-v2
        entry_point=entry_point
    )

    if hasattr(config, 'realRobot'):
        if config.realRobot:
            entry_point = 'Envs.' + env_folder + '.' + ENV + '.ROS' + '.ROS_RL_env:ROSRLEnv'
            register(
                id=ENV + '-ROS-RL-v' + str(model_number),
                entry_point=entry_point
            )

            entry_point = 'Envs.' + env_folder + '.' + ENV + '.ROS' + '.ROS_pretext_env:ROSPretextEnv'
            register(
                id=ENV + '-ROS-pretext-v' + str(model_number),
                entry_point=entry_point
            )
