import numpy as np
import time
from roboimi.envs.double_base import DualDianaMed
# import cv2
import roboimi.utils.KDL_utils.transform as T

import copy as cp


class DualDianaMed_Pos_Ctrl(DualDianaMed):
    """

    :param is_render: Choose if use the renderer to render the scene or not.
    :param renderer: Choose official renderer with "viewer",
            another renderer with "mujoco_viewer"
    :param control_freq: Upper-layer control frequency.
            Note that high frequency will cause high time-lag.
    :param cam_view: Choice of Camera view.
    """

    def __init__(self,
                 robot=None,
                 is_render=False,
                 renderer="viewer",
                 control_freq=200,
                 is_interpolate=True,
                 cam_view=None
                 ):
        super().__init__(
            robot=robot,
            is_render=is_render,
            renderer=renderer,
            control_freq=control_freq,
            is_interpolate=is_interpolate,
            cam_view=cam_view
        )
        
        self.max_reward = 4
        
        self.cam_start()

        
    def step(self,action=np.zeros(16)):
        action_left  = self.ik_solve(action[:3],action[3:7],self.arm_left)
        action_right = self.ik_solve(action[7:10],action[10:14],self.arm_right)
        action = np.hstack((action_left,action_right,action[14:]))
        super().step(action)
        self.rew = self._get_reward()
        
    def step_jnt(self,action):
        super().step(action)

    def ik_solve(self,pos_goal,quat_goal,Arm):
        pos_bias,mat_bias = Arm.get_Arm_bias_frame()
        pos_bias = np.array([pos_bias[0],pos_bias[1],pos_bias[2]])
        mat_bias = np.array([[mat_bias[0],mat_bias[1],mat_bias[2]],
                            [mat_bias[3],mat_bias[4],mat_bias[5]],
                            [mat_bias[6],mat_bias[7],mat_bias[8]]])
        mat_bias = np.linalg.inv(mat_bias)
        mat_goal = T.quat2Mat(quat_goal)
        p_goal = mat_bias@(pos_goal-pos_bias)
        mat_goal = mat_bias@mat_goal
        return Arm.kdl_solver.ikSolver(p_goal, mat_goal, Arm.arm_qpos)

    def reset(self,box_pos):
        
        self.mj_data.joint('red_box_joint').qpos[0] = box_pos[0] 
        self.mj_data.joint('red_box_joint').qpos[1] = box_pos[1]
        self.mj_data.joint('red_box_joint').qpos[2] = box_pos[2]
        self.mj_data.joint('red_box_joint').qpos[3] = 1.0
        self.mj_data.joint('red_box_joint').qpos[4] = 0.0
        self.mj_data.joint('red_box_joint').qpos[5] = 0.0
        self.mj_data.joint('red_box_joint').qpos[6] = 0.0
        super().reset()
        self.top = None
        self.angle = None
        self.r_vis = None
        self.front = None
        self.cam_flage = True
        t=0
        while self.cam_flage:
            if(type(self.top)==type(None) 
               or type(self.angle)==type(None) 
               or type(self.r_vis)==type(None)
               or type(self.front)==type(None)):
                time.sleep(0.001)
                t+=1
            else:
                self.cam_flage=False
        
        

    def preStep(self, action):
        if isinstance(action,np.ndarray) and len(action)==16:
            super().preStep(action)
        else:
            raise AttributeError("size should be 16")

    def get_env_state(self):
        box_pose = np.zeros(3)
        for i in range(3):
            box_pose[i] = cp.deepcopy(self.mj_data.joint('red_box_joint').qpos[i])
        return box_pose
    

    def _get_reward(self):
        all_contact_pairs = []
        for collision_num in range(self.mj_data.ncon):
            geom1 = self.mj_data.contact[collision_num].geom1
            geom2 = self.mj_data.contact[collision_num].geom2
            name_geom_1 = self.getID2Name('geom',geom1)
            name_geom_2 = self.getID2Name('geom',geom2)
            contact_pair = (name_geom_1, name_geom_2)
            all_contact_pairs.append(contact_pair)

        touch_table = ("l_fingertip_g0_right", "table") in all_contact_pairs or\
                      ("l_fingertip_g0_right", "table") in all_contact_pairs
        box_touch_table = ("red_box", "table") in all_contact_pairs
        touch_left_gripper = ("red_box", "r_finger_left") in all_contact_pairs or\
                             ("red_box", "l_finger_left") in all_contact_pairs
        touch_right_gripper = ("red_box", "r_finger_right") in all_contact_pairs or\
                              ("red_box", "l_finger_right") in all_contact_pairs

        reward = 0
        if touch_right_gripper and not touch_table:
            reward = 1
        if touch_right_gripper and not box_touch_table: 
            reward = 2
        if touch_left_gripper: # attempted transfer
            reward = 3
        if touch_left_gripper and not box_touch_table: # successful transfer
            reward = 4
        return reward
        

def make_sim_env(task_name):
    if 'sim_transfer' in task_name:
        from roboimi.assets.robots.diana_med import BiDianaMed
        env = DualDianaMed_Pos_Ctrl(
            robot=BiDianaMed(),
            is_render=True,
            control_freq=30,
            is_interpolate=True,
            cam_view='angle'
        )
        return env
    else:
        raise NotImplementedError





if __name__ == "__main__":
    task_name = 'sim_transfer'
    env = make_sim_env(task_name)
    env.reset()
    action_l = np.array([-0.20,  1.2,  1.1,
                          1.0, 0.0, 0.0, 0.0])  # 双臂直角
    action_r = np.array([0.20,  1.2,  1.1,
                          1.0, 0.0, 0.0, 0.0])
    action_gripper = np.array([-100,-100])
    action = np.hstack((action_l,action_r,action_gripper))

    for t in range(int(10000)):
        # print("main step ", t)
        env.step(action)
        if env.is_render:
            env.render()
        