# Path: osim/env/utils/mygym.py
# def convert_to_gym(space):
#     return spaces.Box(np.array(space[0]), np.array(space[1]) )
# 
# Path: osim/env/osim.py
# class OsimEnv(gym.Env):
#     action_space = None
#     observation_space = None
#     osim_model = None
#     istep = 0
#     verbose = False
# 
#     visualize = False
#     spec = None
#     time_limit = 1e10
# 
#     prev_state_desc = None
# 
#     model_path = None # os.path.join(os.path.dirname(__file__), '../models/MODEL_NAME.osim')    
# 
#     metadata = {
#         'render.modes': ['human'],
#         'video.frames_per_second' : None
#     }
# 
#     def get_reward(self):
#         raise NotImplementedError
# 
#     def is_done(self):
#         return False
# 
#     def __init__(self, visualize = True, integrator_accuracy = 5e-5):
#         self.visualize = visualize
#         self.integrator_accuracy = integrator_accuracy
#         self.load_model()
# 
#     def load_model(self, model_path = None):
#         if model_path:
#             self.model_path = model_path
#             
#         self.osim_model = OsimModel(self.model_path, self.visualize, integrator_accuracy = self.integrator_accuracy)
# 
#         # Create specs, action and observation spaces mocks for compatibility with OpenAI gym
#         self.spec = Spec()
#         self.spec.timestep_limit = self.time_limit
# 
#         self.action_space = ( [0.0] * self.osim_model.get_action_space_size(), [1.0] * self.osim_model.get_action_space_size() )
# #        self.observation_space = ( [-math.pi*100] * self.get_observation_space_size(), [math.pi*100] * self.get_observation_space_s
#         self.observation_space = ( [0] * self.get_observation_space_size(), [0] * self.get_observation_space_size() )
#         
#         self.action_space = convert_to_gym(self.action_space)
#         self.observation_space = convert_to_gym(self.observation_space)
# 
#     def get_state_desc(self):
#         return self.osim_model.get_state_desc()
# 
#     def get_prev_state_desc(self):
#         return self.prev_state_desc
# 
#     def get_observation(self):
#         # This one will normally be overwrtitten by the environments
#         # In particular, for the gym we want a vector and not a dictionary
#         return self.osim_model.get_state_desc()
# 
#     def get_observation_dict(self):
#         return self.osim_model.get_state_desc()
# 
#     def get_observation_space_size(self):
#         return 0
# 
#     def get_action_space_size(self):
#         return self.osim_model.get_action_space_size()
# 
#     def reset(self, project=True, obs_as_dict=True):
#         self.osim_model.reset()
# 
#         if not project:
#             return self.get_state_desc()
#         if obs_as_dict:
#             return self.get_observation_dict()
#         return self.get_observation()
# 
#     def step(self, action, project=True, obs_as_dict=True):
#         self.prev_state_desc = self.get_state_desc()        
#         self.osim_model.actuate(action)
#         self.osim_model.integrate()
# 
#         if project:
#             if obs_as_dict:
#                 obs = self.get_observation_dict()
#             else:
#                 obs = self.get_observation()
#         else:
#             obs = self.get_state_desc()
#             
#         return [ obs, self.get_reward(), self.is_done() or (self.osim_model.istep >= self.spec.timestep_limit), {} ]
# 
#     def render(self, mode='human', close=False):
#         return


import math
import numpy as np
import os
import gym
import opensim
import random
from .utils.mygym import convert_to_gym
from .osim import OsimEnv
    model_path = os.path.join(os.path.dirname(__file__), '../models/arm2dof6musc.osim')    
    time_limit = 200
    target_x = 0
    target_y = 0

    def get_observation(self):
        state_desc = self.get_state_desc()

        res = [self.target_x, self.target_y]

        # for body_part in ["r_humerus", "r_ulna_radius_hand"]:
        #     res += state_desc["body_pos"][body_part][0:2]
        #     res += state_desc["body_vel"][body_part][0:2]
        #     res += state_desc["body_acc"][body_part][0:2]
        #     res += state_desc["body_pos_rot"][body_part][2:]
        #     res += state_desc["body_vel_rot"][body_part][2:]
        #     res += state_desc["body_acc_rot"][body_part][2:]

        for joint in ["r_shoulder","r_elbow",]:
            res += state_desc["joint_pos"][joint]
            res += state_desc["joint_vel"][joint]
            res += state_desc["joint_acc"][joint]

        for muscle in sorted(state_desc["muscles"].keys()):
            res += [state_desc["muscles"][muscle]["activation"]]
            # res += [state_desc["muscles"][muscle]["fiber_length"]]
            # res += [state_desc["muscles"][muscle]["fiber_velocity"]]

        res += state_desc["markers"]["r_radius_styloid"]["pos"][:2]


# Path: osim/env/utils/mygym.py
# def convert_to_gym(space):
#     return spaces.Box(np.array(space[0]), np.array(space[1]) )
# 
# Path: osim/env/osim.py
# class OsimEnv(gym.Env):
#     action_space = None
#     observation_space = None
#     osim_model = None
#     istep = 0
#     verbose = False
# 
#     visualize = False
#     spec = None
#     time_limit = 1e10
# 
#     prev_state_desc = None
# 
#     model_path = None # os.path.join(os.path.dirname(__file__), '../models/MODEL_NAME.osim')    
# 
#     metadata = {
#         'render.modes': ['human'],
#         'video.frames_per_second' : None
#     }
# 
#     def get_reward(self):
#         raise NotImplementedError
# 
#     def is_done(self):
#         return False
# 
#     def __init__(self, visualize = True, integrator_accuracy = 5e-5):
#         self.visualize = visualize
#         self.integrator_accuracy = integrator_accuracy
#         self.load_model()
# 
#     def load_model(self, model_path = None):
#         if model_path:
#             self.model_path = model_path
#             
#         self.osim_model = OsimModel(self.model_path, self.visualize, integrator_accuracy = self.integrator_accuracy)
# 
#         # Create specs, action and observation spaces mocks for compatibility with OpenAI gym
#         self.spec = Spec()
#         self.spec.timestep_limit = self.time_limit
# 
#         self.action_space = ( [0.0] * self.osim_model.get_action_space_size(), [1.0] * self.osim_model.get_action_space_size() )
# #        self.observation_space = ( [-math.pi*100] * self.get_observation_space_size(), [math.pi*100] * self.get_observation_space_s
#         self.observation_space = ( [0] * self.get_observation_space_size(), [0] * self.get_observation_space_size() )
#         
#         self.action_space = convert_to_gym(self.action_space)
#         self.observation_space = convert_to_gym(self.observation_space)
# 
#     def get_state_desc(self):
#         return self.osim_model.get_state_desc()
# 
#     def get_prev_state_desc(self):
#         return self.prev_state_desc
# 
#     def get_observation(self):
#         # This one will normally be overwrtitten by the environments
#         # In particular, for the gym we want a vector and not a dictionary
#         return self.osim_model.get_state_desc()
# 
#     def get_observation_dict(self):
#         return self.osim_model.get_state_desc()
# 
#     def get_observation_space_size(self):
#         return 0
# 
#     def get_action_space_size(self):
#         return self.osim_model.get_action_space_size()
# 
#     def reset(self, project=True, obs_as_dict=True):
#         self.osim_model.reset()
# 
#         if not project:
#             return self.get_state_desc()
#         if obs_as_dict:
#             return self.get_observation_dict()
#         return self.get_observation()
# 
#     def step(self, action, project=True, obs_as_dict=True):
#         self.prev_state_desc = self.get_state_desc()        
#         self.osim_model.actuate(action)
#         self.osim_model.integrate()
# 
#         if project:
#             if obs_as_dict:
#                 obs = self.get_observation_dict()
#             else:
#                 obs = self.get_observation()
#         else:
#             obs = self.get_state_desc()
#             
#         return [ obs, self.get_reward(), self.is_done() or (self.osim_model.istep >= self.spec.timestep_limit), {} ]
# 
#     def render(self, mode='human', close=False):
#         return


# Path: osim/env/arm.py
import math
import numpy as np
import os
import gym
import opensim
import random
from .utils.mygym import convert_to_gym
from .osim import OsimEnv

    model_path = os.path.join(os.path.dirname(__file__), '../models/arm2dof6musc.osim')    
    time_limit = 200
    target_x = 0
    target_y = 0

    def get_observation(self):
        state_desc = self.get_state_desc()

        res = [self.target_x, self.target_y]

        # for body_part in ["r_humerus", "r_ulna_radius_hand"]:
        #     res += state_desc["body_pos"][body_part][0:2]
        #     res += state_desc["body_vel"][body_part][0:2]
        #     res += state_desc["body_acc"][body_part][0:2]
        #     res += state_desc["body_pos_rot"][body_part][2:]
        #     res += state_desc["body_vel_rot"][body_part][2:]
        #     res += state_desc["body_acc_rot"][body_part][2:]

        for joint in ["r_shoulder","r_elbow",]:
            res += state_desc["joint_pos"][joint]
            res += state_desc["joint_vel"][joint]
            res += state_desc["joint_acc"][joint]

        for muscle in sorted(state_desc["muscles"].keys()):
            res += [state_desc["muscles"][muscle]["activation"]]
            # res += [state_desc["muscles"][muscle]["fiber_length"]]
            # res += [state_desc["muscles"][muscle]["fiber_velocity"]]

        res += state_desc["markers"]["r_radius_styloid"]["pos"][:2]


        return res