# Path: util/logger.py
# class _MyFormatter(logging.Formatter):
#     def format(self, record):
# def set_file_handler(path=None, prefix='', time_str=''):
# _LOGGING_METHOD = ['info', 'warning', 'error', 'critical',
#                    'warn', 'exception', 'debug']
# 
# Path: environments/register.py
# MAX_EPISODE_STEPS_DICT = {
#     'Centipede': 1000,
#     'CpCentipede': 1000,
#     'Snake': 1000,
#     'Reacher': 50,
# }
# REWARD_THRESHOLD = {
#     'Centipede': 6000.0,
#     'CpCentipede': 6000.0,
#     'Snake': 360.0,
#     'Reacher': -3.75,
# }
# MULTI_TASK_DICT = {
#     'MultiWalkers-v1':
#         ['WalkersHopper-v1', 'WalkersHalfhumanoid-v1', 'WalkersHalfcheetah-v1',
#          'WalkersFullcheetah-v1', 'WalkersOstrich-v1'],
#     # just for implementation, only one agent will be run
#     'MultiWalkers2Kangaroo-v1':
#         ['WalkersHopper-v1', 'WalkersHalfhumanoid-v1', 'WalkersHalfcheetah-v1',
#          'WalkersFullcheetah-v1', 'WalkersKangaroo-v1'],
# }
# NUM_ROBUSTNESS_AGENTS = 5
# ROBUSTNESS_TASK_DICT = {}
#     JOINT_KEY = ['qpos', 'qvel', 'qfrc_constr', 'qfrc_act']
#     BODY_KEY = ['cinert', 'cvel', 'cfrc']
#     ROOT_OB_SIZE = {
#         'qpos': {'free': 7, 'hinge': 1, 'slide': 1},
#         'qvel': {'free': 6, 'hinge': 1, 'slide': 1},
#         'qfrc_act': {'free': 6, 'hinge': 1, 'slide': 1},
#         'qfrc_constr': {'free': 6, 'hinge': 1, 'slide': 1}
#     }
#     SYMMETRY_MAP = {'Humanoid-v1': 2,
#                     'HumanoidStandup-v1': 2,
#                     'HalfCheetah-v1': 1,
#                     'Hopper-v1': 1,
#                     'Walker2d-v1': 1,
#                     'AntS-v1': 2,
#                     'Swimmer-v1': 2,
# 
#                     'WalkersHopper-v1': 1,
#                     'WalkersHalfhumanoid-v1': 1,
#                     'WalkersHalfcheetah-v1': 1,
#                     'WalkersFullcheetah-v1': 1,
#                     'WalkersOstrich-v1': 1,
#                     'WalkersKangaroo-v1': 1}
#     XML_DICT = {'Humanoid-v1': 'humanoid.xml',
#                 'HumanoidStandup-v1': 'humanoid.xml',
#                 'HalfCheetah-v1': 'half_cheetah.xml',
#                 'Hopper-v1': 'hopper.xml',
#                 'Walker2d-v1': 'walker2d.xml',
#                 'AntS-v1': 'ant.xml',
#                 'Swimmer-v1': 'SnakeThree.xml',
# 
#                 'WalkersHopper-v1': 'WalkersHopper.xml',
#                 'WalkersHalfhumanoid-v1': 'WalkersHalfhumanoid.xml',
#                 'WalkersHalfcheetah-v1': 'WalkersHalfcheetah.xml',
#                 'WalkersFullcheetah-v1': 'WalkersFullcheetah.xml',
#                 'WalkersOstrich-v1': 'WalkersOstrich.xml',
#                 'WalkersKangaroo-v1': 'WalkersKangaroo.xml'}
#     OB_MAP = {
#         'Humanoid-v1':
#             ['qpos', 'qvel', 'cinert', 'cvel', 'qfrc_act', 'cfrc'],
#         'HumanoidStandup-v1':
#             ['qpos', 'qvel', 'cinert', 'cvel', 'qfrc_act', 'cfrc'],
#         'HalfCheetah-v1': ['qpos', 'qvel'],
#         'Hopper-v1': ['qpos', 'qvel'],
#         'Walker2d-v1': ['qpos', 'qvel'],
#         'AntS-v1': ['qpos', 'qvel', 'cfrc'],
#         'Swimmer-v1': ['qpos', 'qvel'],
# 
#         'WalkersHopper-v1': ['qpos', 'qvel'],
#         'WalkersHalfhumanoid-v1': ['qpos', 'qvel'],
#         'WalkersHalfcheetah-v1': ['qpos', 'qvel'],
#         'WalkersFullcheetah-v1': ['qpos', 'qvel'],
#         'WalkersOstrich-v1': ['qpos', 'qvel'],
#         'WalkersKangaroo-v1': ['qpos', 'qvel']
#     }
# def get_name_list():
# def get_mujoco_model_settings():


import init_path
import os
import numpy as np
from bs4 import BeautifulSoup as bs
from util import logger
from environments import register
'''
    Definition of nodes:
    @root:
        The 'root' type is the combination of the top level 'body' node and
        the top level free 'joint' (two nodes combined)
        Also, additional input will be assigned to the root node
        (e.g. the postion of the targer).

        For different tasks, we should have different MLP for each root.

    @geom, @body, @joint:
        The structure defined in the xml files. Ideally, the MLP for input,
        propogation, and output could be shared among different models.
'''
_NODE_TYPE = ['root', 'joint', 'body']  # geom is removed
EDGE_TYPE = {'self_loop': 0, 'root-root': 0,  # root-root is loop, also 0
             'joint-joint': 1, 'geom-geom': 2, 'body-body': 3, 'tendon': 4,
             'joint-geom': 5, 'geom-joint': -5,  # pc-relationship
             'joint-body': 6, 'body-joint': -6,
             'body-geom': 7, 'geom-body': -7,
             'root-geom': 8, 'geom-root': -8,
             'root-joint': 9, 'joint-root': -9,
             'root-body': 10, 'body-root': -10}

SYMMETRY_MAP, XML_DICT, OB_MAP, JOINT_KEY, ROOT_OB_SIZE, BODY_KEY = \
    register.get_mujoco_model_settings()

# development status
ALLOWED_JOINT_TYPE = ['hinge', 'free', 'slide']


# Path: util/logger.py
# class _MyFormatter(logging.Formatter):
#     def format(self, record):
# def set_file_handler(path=None, prefix='', time_str=''):
# _LOGGING_METHOD = ['info', 'warning', 'error', 'critical',
#                    'warn', 'exception', 'debug']
# 
# Path: environments/register.py
# MAX_EPISODE_STEPS_DICT = {
#     'Centipede': 1000,
#     'CpCentipede': 1000,
#     'Snake': 1000,
#     'Reacher': 50,
# }
# REWARD_THRESHOLD = {
#     'Centipede': 6000.0,
#     'CpCentipede': 6000.0,
#     'Snake': 360.0,
#     'Reacher': -3.75,
# }
# MULTI_TASK_DICT = {
#     'MultiWalkers-v1':
#         ['WalkersHopper-v1', 'WalkersHalfhumanoid-v1', 'WalkersHalfcheetah-v1',
#          'WalkersFullcheetah-v1', 'WalkersOstrich-v1'],
#     # just for implementation, only one agent will be run
#     'MultiWalkers2Kangaroo-v1':
#         ['WalkersHopper-v1', 'WalkersHalfhumanoid-v1', 'WalkersHalfcheetah-v1',
#          'WalkersFullcheetah-v1', 'WalkersKangaroo-v1'],
# }
# NUM_ROBUSTNESS_AGENTS = 5
# ROBUSTNESS_TASK_DICT = {}
#     JOINT_KEY = ['qpos', 'qvel', 'qfrc_constr', 'qfrc_act']
#     BODY_KEY = ['cinert', 'cvel', 'cfrc']
#     ROOT_OB_SIZE = {
#         'qpos': {'free': 7, 'hinge': 1, 'slide': 1},
#         'qvel': {'free': 6, 'hinge': 1, 'slide': 1},
#         'qfrc_act': {'free': 6, 'hinge': 1, 'slide': 1},
#         'qfrc_constr': {'free': 6, 'hinge': 1, 'slide': 1}
#     }
#     SYMMETRY_MAP = {'Humanoid-v1': 2,
#                     'HumanoidStandup-v1': 2,
#                     'HalfCheetah-v1': 1,
#                     'Hopper-v1': 1,
#                     'Walker2d-v1': 1,
#                     'AntS-v1': 2,
#                     'Swimmer-v1': 2,
# 
#                     'WalkersHopper-v1': 1,
#                     'WalkersHalfhumanoid-v1': 1,
#                     'WalkersHalfcheetah-v1': 1,
#                     'WalkersFullcheetah-v1': 1,
#                     'WalkersOstrich-v1': 1,
#                     'WalkersKangaroo-v1': 1}
#     XML_DICT = {'Humanoid-v1': 'humanoid.xml',
#                 'HumanoidStandup-v1': 'humanoid.xml',
#                 'HalfCheetah-v1': 'half_cheetah.xml',
#                 'Hopper-v1': 'hopper.xml',
#                 'Walker2d-v1': 'walker2d.xml',
#                 'AntS-v1': 'ant.xml',
#                 'Swimmer-v1': 'SnakeThree.xml',
# 
#                 'WalkersHopper-v1': 'WalkersHopper.xml',
#                 'WalkersHalfhumanoid-v1': 'WalkersHalfhumanoid.xml',
#                 'WalkersHalfcheetah-v1': 'WalkersHalfcheetah.xml',
#                 'WalkersFullcheetah-v1': 'WalkersFullcheetah.xml',
#                 'WalkersOstrich-v1': 'WalkersOstrich.xml',
#                 'WalkersKangaroo-v1': 'WalkersKangaroo.xml'}
#     OB_MAP = {
#         'Humanoid-v1':
#             ['qpos', 'qvel', 'cinert', 'cvel', 'qfrc_act', 'cfrc'],
#         'HumanoidStandup-v1':
#             ['qpos', 'qvel', 'cinert', 'cvel', 'qfrc_act', 'cfrc'],
#         'HalfCheetah-v1': ['qpos', 'qvel'],
#         'Hopper-v1': ['qpos', 'qvel'],
#         'Walker2d-v1': ['qpos', 'qvel'],
#         'AntS-v1': ['qpos', 'qvel', 'cfrc'],
#         'Swimmer-v1': ['qpos', 'qvel'],
# 
#         'WalkersHopper-v1': ['qpos', 'qvel'],
#         'WalkersHalfhumanoid-v1': ['qpos', 'qvel'],
#         'WalkersHalfcheetah-v1': ['qpos', 'qvel'],
#         'WalkersFullcheetah-v1': ['qpos', 'qvel'],
#         'WalkersOstrich-v1': ['qpos', 'qvel'],
#         'WalkersKangaroo-v1': ['qpos', 'qvel']
#     }
# def get_name_list():
# def get_mujoco_model_settings():


# Path: graph_util/mujoco_parser.py
import init_path
import os
import numpy as np
from bs4 import BeautifulSoup as bs
from util import logger
from environments import register

'''
    Definition of nodes:
    @root:
        The 'root' type is the combination of the top level 'body' node and
        the top level free 'joint' (two nodes combined)
        Also, additional input will be assigned to the root node
        (e.g. the postion of the targer).

        For different tasks, we should have different MLP for each root.

    @geom, @body, @joint:
        The structure defined in the xml files. Ideally, the MLP for input,
        propogation, and output could be shared among different models.
'''
_NODE_TYPE = ['root', 'joint', 'body']  # geom is removed
EDGE_TYPE = {'self_loop': 0, 'root-root': 0,  # root-root is loop, also 0
             'joint-joint': 1, 'geom-geom': 2, 'body-body': 3, 'tendon': 4,
             'joint-geom': 5, 'geom-joint': -5,  # pc-relationship
             'joint-body': 6, 'body-joint': -6,
             'body-geom': 7, 'geom-body': -7,
             'root-geom': 8, 'geom-root': -8,
             'root-joint': 9, 'joint-root': -9,
             'root-body': 10, 'body-root': -10}

SYMMETRY_MAP, XML_DICT, OB_MAP, JOINT_KEY, ROOT_OB_SIZE, BODY_KEY = \
    register.get_mujoco_model_settings()

# development status
ALLOWED_JOINT_TYPE = ['hinge', 'free', 'slide']


MULTI_TASK_DICT = register.MULTI_TASK_DICT