#Minimal implementation of the RL_Framework_Mujoco environment for CMA-ES and IK Optimization

#for finding the inital pose of the musculoskeletal model

#and for trajectory visualization

from gym import error, spaces

from gym.utils import seeding

import numpy as np

from os import path

import gym

import pickle

import numpy as np

from gym import utils

from . import sensory_feedback_specs

from . import kinematics_preprocessing_specs

import ipdb

try:

    import mujoco_py

except ImportError as e:

    raise error.DependencyNotInstalled("{}. (HINT: you need to install mujoco_py, and also perform the setup instructions here: https://github.com/openai/mujoco-py/.)".format(e))

DEFAULT_SIZE = 500

class MujocoEnv(gym.Env):

    """Superclass for all MuJoCo environments.

    """

    def __init__(self, model_path, condition_to_sim, cond_tpoint, args):

        self.model_path = model_path

        self.initial_pose_path = args.initial_pose_path

        self.kinematics_path = args.kinematics_path

        self.cond_to_sim = condition_to_sim

        self.cond_tpoint = cond_tpoint

         #setup the model

        self.model = mujoco_py.load_model_from_path(model_path)

        #setup the simulator and data object

        self.sim = mujoco_py.MjSim(self.model)

        self.data = self.sim.data

        #Save all the sensory feedback specs for use in the later functions

        self.sfs_proprioceptive_feedback = args.proprioceptive_feedback

        self.sfs_muscle_forces = args.muscle_forces

        self.sfs_joint_feedback = args.joint_feedback

        self.sfs_visual_feedback = args.visual_feedback

        self.sfs_visual_feedback_bodies = args.visual_feedback_bodies

        self.sfs_visual_distance_bodies = args.visual_distance_bodies

        self.sfs_visual_velocity = args.visual_velocity

        self.sfs_sensory_delay_timepoints = args.sensory_delay_timepoints

        #Load the kin_train

        # Load the experimental kinematics x and y coordinates from the data

        with open(self.kinematics_path + '/kinematics.pkl', 'rb') as f:

            kin_train_test = pickle.load(f)

        kin_train = kin_train_test['train'] #[num_conds][num_targets, num_coords, timepoints]

        #Randomly sample the number of targets from a condition = 0

        num_targets = kin_train[0].shape[0]

        #Find the qpos corresponding to the musculoskeletal model and targets

        self.qpos_idx_musculo = np.array(list(range(0, self.model.nq)))

        self.qpos_idx_targets = []

        for musculo_targets in kinematics_preprocessing_specs.musculo_target_joints:

            joint_idx = self.model.get_joint_qpos_addr(musculo_targets)

            self.qpos_idx_targets.append(joint_idx)

        #Delete the corresponding index from the qpos_idx_musculo

        self.qpos_idx_musculo = np.delete(self.qpos_idx_musculo, self.qpos_idx_targets).tolist()

        #Set the simulation timestep

        if args.sim_dt != 0:

            self.model.opt.timestep = args.sim_dt

        self.n_fixedsteps = args.n_fixedsteps

        self.radius = args.trajectory_scaling

        self.center = args.center

        #Kinematics preprocessing for training and testing kinematics

        #Preprocess training kinematics

        for i_target in range(kin_train[0].shape[0]):

            for i_cond in range(len(kin_train)):

                for i_coord in range(kin_train[i_cond].shape[1]):

                    kin_train[i_cond][i_target, i_coord, :] = kin_train[i_cond][i_target, i_coord, :] / self.radius[i_target]

                    kin_train[i_cond][i_target, i_coord, :] = kin_train[i_cond][i_target, i_coord, :] + self.center[i_target][i_coord]

        self.kin_train = kin_train

        self.kin_to_sim = self.kin_train

        self.current_cond_to_sim = self.cond_to_sim

        self.upd_theta(self.cond_tpoint)

        self.viewer = None 

        self._viewers = {}

        self.metadata = {

            'render.modes': ['human', 'rgb_array', 'depth_array'],

            'video.frames_per_second': int(np.round(1.0 / self.dt))

        }

        #init_qpos and init_qvel do not contain target qpos

        #and contain only the musculo qpos

        #If the initial qpos and qvel are provided by the user

        if path.isfile(self.initial_pose_path + '/qpos.npy'):

            init_qpos = np.load(self.initial_pose_path + '/qpos.npy')

            init_qvel = np.load(self.initial_pose_path + '/qvel.npy')

        #else use the default initial pose of xml model

        else:

            init_qpos = self.sim.data.qpos.flat.copy()[self.qpos_idx_musculo]

            init_qvel = self.sim.data.qvel.flat.copy()[self.qpos_idx_musculo]

        #Get the qpos of musclo + targets

        musculo_qpos = self.sim.data.qpos.flat.copy() 

        musculo_qvel = self.sim.data.qvel.flat.copy()

        #Set the musculo part to the saved initial qpos and qvel

        musculo_qpos[self.qpos_idx_musculo] = init_qpos

        self.init_qpos = musculo_qpos

        #Set the initial state to init_qpos

        #Set the state to the initial pose

        self.set_state(self.init_qpos)

        if self.sfs_visual_feedback == True:

            #Save the xpos of the musculo bodies for visual vels

            if len(self.sfs_visual_velocity) != 0:

                self.prev_body_xpos = []

                for musculo_body in self.sfs_visual_velocity:

                    body_xpos = self.sim.data.get_body_xpos(musculo_body)

                    self.prev_body_xpos = [*self.prev_body_xpos, *body_xpos]

    #Return the qpos of only the musculoskeletal bodies and not the targets

    def get_musculo_state(self):

        qpos_all = self.sim.data.qpos.flat.copy()

        qpos_musculo = qpos_all[self.qpos_idx_musculo]

        return qpos_musculo

    def set_state(self, qpos):

        assert qpos.shape == (self.model.nq, )

        old_state= self.sim.get_state()

        new_state= mujoco_py.MjSimState(old_state.time, qpos, qpos*0,

                                        old_state.act, old_state.udd_state)

        self.sim.set_state(new_state)

        self.sim.forward()

    def set_state_musculo(self, qpos):

        old_state= self.sim.get_state()

        qpos_all = self.sim.data.qpos.flat.copy()

        qpos_all[self.qpos_idx_musculo] = qpos

        assert qpos_all.shape == (self.model.nq, )

        new_state= mujoco_py.MjSimState(old_state.time, qpos_all, qpos_all*0,

                                        old_state.act, old_state.udd_state)

        self.sim.set_state(new_state)

        self.sim.forward()

    def get_obs_musculo_bodies(self):

        #Returns the current xyz coords of the musculo bodies to be tracked

        musculo_body_state = []

        for musculo_body in kinematics_preprocessing_specs.musculo_tracking:

            current_xyz_coord = self.sim.data.get_body_xpos(musculo_body[0]).flat.copy() 

            musculo_body_state.append(current_xyz_coord)

        return np.array(musculo_body_state)  #[n_musculo_bodies, 3]

    def get_obs_targets(self):

        #Returns the current xyz coords of the targets to be tracked

        musculo_target_state = []

        for musculo_body in kinematics_preprocessing_specs.musculo_tracking:

            current_xyz_coord = self.sim.data.get_body_xpos(musculo_body[1]).flat.copy() 

            musculo_target_state.append(current_xyz_coord)

        return np.array(musculo_target_state)  #[n_musculo_targets, 3]

    def set_cond_to_simulate(self, i_condition, cond_timepoint):

        #update the state variables

        self.cond_to_sim = i_condition

        self.cond_tpoint = cond_timepoint

        self.current_cond_to_sim = self.cond_to_sim

        self.upd_theta(self.cond_tpoint)

    @property

    def dt(self):

        return self.model.opt.timestep 

    def render(self,

               mode='human',

               width=DEFAULT_SIZE,

               height=DEFAULT_SIZE,

               camera_id=0,

               camera_name=None):

        if mode == 'rgb_array' or mode == 'depth_array':

            if camera_id is not None and camera_name is not None:

                raise ValueError("Both `camera_id` and `camera_name` cannot be"

                                 " specified at the same time.")

            no_camera_specified = camera_name is None and camera_id is None

            if no_camera_specified:

                camera_name = 'track'

            if camera_id is None and camera_name in self.model._camera_name2id:

                camera_id = self.model.camera_name2id(camera_name)

            self._get_viewer(mode).render(width, height, camera_id=camera_id)

        if mode == 'rgb_array':

            # window size used for old mujoco-py:

            data = self._get_viewer(mode).read_pixels(width, height, depth=False)

            # original image is upside-down, so flip it

            return data[::-1, :, :]

        elif mode == 'depth_array':

            self._get_viewer(mode).render(width, height)

            # window size used for old mujoco-py:

            # Extract depth part of the read_pixels() tuple

            data = self._get_viewer(mode).read_pixels(width, height, depth=True)[1]

            # original image is upside-down, so flip it

            return data[::-1, :]

        elif mode == 'human':

            self._get_viewer(mode).render()

    def close(self):

        if self.viewer is not None:

            # self.viewer.finish()

            self.viewer = None

            self._viewers = {}

    def _get_viewer(self, mode):

        self.viewer = self._viewers.get(mode)

        if self.viewer is None:

            if mode == 'human':

                self.viewer = mujoco_py.MjViewer(self.sim)

            elif mode == 'rgb_array' or mode == 'depth_array':

                self.viewer = mujoco_py.MjRenderContextOffscreen(self.sim, -1)

            self.viewer_setup()

            self._viewers[mode] = self.viewer

        return self.viewer

    def get_body_com(self, body_name):

        return self.data.get_body_xpos(body_name).copy()

class Muscle_Env(MujocoEnv):

    def __init__(self, model_path, condition_to_sim, cond_tpoint, args):

        MujocoEnv.__init__(self, model_path, condition_to_sim, cond_tpoint, args)

    def viewer_setup(self):

        self.viewer.cam.trackbodyid = 0

    def upd_theta(self, cond_timepoint):

        coords_to_sim = self.kin_to_sim[self.current_cond_to_sim] #[num_targets, num_coords, timepoints]

        assert cond_timepoint < coords_to_sim.shape[-1]

        crnt_state = self.sim.get_state()

        for i_target in range(self.kin_to_sim[self.current_cond_to_sim].shape[0]):

            if kinematics_preprocessing_specs.xyz_target[i_target][0]:

                x_joint_idx= self.model.get_joint_qpos_addr(f"box:x{i_target}")

                crnt_state.qpos[x_joint_idx] = coords_to_sim[i_target, 0, cond_timepoint]

            if kinematics_preprocessing_specs.xyz_target[i_target][1]:

                y_joint_idx= self.model.get_joint_qpos_addr(f"box:y{i_target}")

                crnt_state.qpos[y_joint_idx] = coords_to_sim[i_target, kinematics_preprocessing_specs.xyz_target[i_target][0], cond_timepoint]

            if kinematics_preprocessing_specs.xyz_target[i_target][2]:

                z_joint_idx= self.model.get_joint_qpos_addr(f"box:z{i_target}")

                crnt_state.qpos[z_joint_idx] = coords_to_sim[i_target, kinematics_preprocessing_specs.xyz_target[i_target][0] + kinematics_preprocessing_specs.xyz_target[i_target][1], cond_timepoint]

        #Now set the state

        self.set_state(crnt_state.qpos)

    def _get_obs(self):

        sensory_feedback = []

        if self.sfs_proprioceptive_feedback == True:

            muscle_lens = self.sim.data.actuator_length.flat.copy()

            muscle_vels = self.sim.data.actuator_velocity.flat.copy()

            #process through the given function for muscle lens and muscle vels

            muscle_lens, muscle_vels = sensory_feedback_specs.process_proprioceptive(muscle_lens, muscle_vels)

            sensory_feedback = [*sensory_feedback, *muscle_lens, *muscle_vels]

        if self.sfs_muscle_forces == True:

            actuator_forces = self.sim.data.qfrc_actuator.flat.copy()

            #process

            actuator_forces = sensory_feedback_specs.process_muscle_forces(actuator_forces)

            sensory_feedback = [*sensory_feedback, *actuator_forces]

        if self.sfs_joint_feedback == True:

            sensory_qpos = self.sim.data.qpos.flat.copy()

            sensory_qvel = self.sim.data.qvel.flat.copy()

            sensory_qpos, sensory_qvel = sensory_feedback_specs.process_joint_feedback(sensory_qpos, sensory_qvel)

            sensory_feedback = [*sensory_feedback, *sensory_qpos, *sensory_qvel]

        if self.sfs_visual_feedback == True:

            #Check if the user specified the musculo bodies to be included

            assert len(self.sfs_visual_feedback_bodies) != 0

            visual_xyz_coords = []

            for musculo_body in self.sfs_visual_feedback_bodies:

                visual_xyz_coords = [*visual_xyz_coords, *self.sim.data.get_body_xpos(musculo_body)]

            visual_xyz_coords = sensory_feedback_specs.process_visual_position(visual_xyz_coords)

            sensory_feedback = [*sensory_feedback, *visual_xyz_coords]

        if len(self.sfs_visual_distance_bodies) != 0:

            visual_xyz_distance = []

            for musculo_tuple in self.sfs_visual_distance_bodies:

                body0_xyz = self.sim.data.get_body_xpos(musculo_tuple[0])

                body1_xyz = self.sim.data.get_body_xpos(musculo_tuple[1])

                tuple_dist = np.abs(body0_xyz - body1_xyz).tolist()

                visual_xyz_distance = [*visual_xyz_distance, *tuple_dist]

            #process

            visual_xyz_distance = sensory_feedback_specs.process_visual_distance(visual_xyz_distance)

            sensory_feedback = [*sensory_feedback, *visual_xyz_distance]

        #Save the xpos of the musculo bodies for visual vels

        if len(self.sfs_visual_velocity) != 0:

            #Find the visual vels after the simulation

            current_body_xpos = []

            for musculo_body in self.sfs_visual_velocity:

                body_xpos = self.sim.data.get_body_xpos(musculo_body)

                current_body_xpos = [*current_body_xpos, *body_xpos]

            #Find the velocity

            visual_vels = (np.abs(np.array(self.prev_body_xpos) - np.array(current_body_xpos)) / self.dt).tolist()

            #process visual velocity feedback

            visual_vels = sensory_feedback_specs.process_visual_velocity(visual_vels)

            sensory_feedback= [*sensory_feedback, *visual_vels]

            #update the self.prev_body_xpos

            self.prev_body_xpos = current_body_xpos

        return np.array(sensory_feedback)