import numpy as np
from roboimi.assets.robots.bidiana_med_vs import BiDianaMedVS
from roboimi.envs.visual_servo.double_vs_base import DualVisualServoEnv
from scipy.spatial.transform import Rotation as RR

class RLDualVisualServoEnv(DualVisualServoEnv):
    def __init__(self,
                 robot=None,
                 is_render=False,
                 renderer="viewer",
                 control_freq=200,
                 is_interpolate=False,
                 ):
        super().__init__(
            robot=robot,
            is_render=is_render,
            renderer=renderer,
            control_freq=control_freq,
            is_interpolate=is_interpolate
        )
        self.action_low = -1.0
        self.action_high = 1.0
        self.obs_low =  np.concatenate((np.array([-3.1415]*7),
                                        np.array([-3.0]*7),
                                        np.array([-10.0]*3),
                                        np.array([-1.0]*4)))
        self.obs_high = np.concatenate((np.array([3.1415]*7),
                                        np.array([3.0]*7),
                                        np.array([10.0]*3),
                                        np.array([1.0]*4)))
        self.obs_dim = (21,)
        self.max_episode_steps = 1000
        self.name = "vs_env"


    def preStep(self, action):
        return super().preStep(action)


    def step(self,action):
        super().step(action)
        done = True
        info={}
        return self._get_obs(), self._get_reward(), done, True, info

    # rl env must have
    def reset(self):
        super().reset()
        info = {}
        return self._get_obs(), info
    

    def _get_obs(self):
        '''
        obs   0~6: arm jnt pos
        obs  7~13: arm vel pos
        obs 14~16: arm eef position [x y z]
        obs 17~20: arm eef quaternion [w qx qy qz]
        '''
        obs = np.zeros(21,dtype=np.float32)
        # arm_qpos = np.asarray(self.arm.arm_qpos)
        # print(type(arm_qpos[0]),"\n--",type(arm_qpos[1]))
        obs[:7] = np.asarray(self.arm.arm_qpos,dtype=np.float32)
        obs[7:14] = np.asarray(self.arm.arm_qvel,dtype=np.float32)
        obs[14:17] = self.getBodyPos("0_end_effector")
        obs[17:21] = self.getBodyQuat("0_end_effector")
        return obs
        
        
    def _get_reward(self):
        reward = 0.0
        return reward


if __name__ == "__main__":
    try:
        from roboimi.assets.robots.bidiana_med_vs import BiDianaMedVS
        env = RLDualVisualServoEnv(
            robot=BiDianaMedVS(),
            is_render=True,
            control_freq=200,
            is_interpolate=True,
        )
        env.reset()
        env.cam_start()
        for t in range(int(1e6)):
            # action = np.array([0.63116, -0.09768533, 0.46947228, 1, 1, 0, 0.63116, -0.09768533, 0.46947228, 1, 0, 0])
            # action = np.zeros(12)
            # action[0] = 1.0
            # action[6] = 1.0
            # env.step(action)
            env.step()
            if env.is_render:
                env.render()
    except:
        KeyboardInterrupt