import numpy as np

from data_gen_dependencies.action.base import StageTemplate


class TwistStage(StageTemplate):
    DELTA_DISTANCE = 0.01  # meter
    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), extra_params=None, objects = [], **kwargs):
        super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
        self.passive_obj = objects[passive_obj_id]
        self.active_obj = objects[active_obj_id]
        self.joint_position_threshold = passive_element.get('joint_position_threshold', 0.7)
        if self.joint_position_threshold is None:
            self.joint_position_threshold = 0.7
        self.correspond_joint_id = passive_element.get('correspond_joint_id', None)
        correspond_joint_info = self.passive_obj.joints_info[self.correspond_joint_id]
        self.joint_lower_limit = correspond_joint_info["lower_bound"]
        self.joint_upper_limit = correspond_joint_info["upper_bound"]
        if self.joint_lower_limit is None or self.joint_upper_limit is None:
            self.twist_degree_range = np.pi
        else:
            self.twist_degree_range = abs(self.joint_upper_limit - self.joint_lower_limit)
        self.joint_axis = correspond_joint_info["joint_axis"]
        self.joint_type = correspond_joint_info["joint_type"]
        assert self.joint_type == 'continuous', "joint_type must be continuous for twist action"
        if self.joint_position_threshold is None:
            self.joint_position_threshold = 0.7
        assert self.joint_position_threshold >= 0 and self.joint_position_threshold <= 1 

        self.joint_direction = passive_element.get('joint_direction', 1)
        assert self.joint_direction in [-1, 1]

        self.joint_velocity_threshold = passive_element.get('joint_velocity_threshold', 999)
        vector_direction = passive_element['direction']

        self.extra_params = {} if extra_params is None else extra_params
        self.generate_substage(target_pose, vector_direction)

    def generate_substage(self, target_pose, vector_direction):
        vector_direction = vector_direction / np.linalg.norm(vector_direction)
        delta_twist_pose = self.axis_angle_to_matrix(np.asarray([0,0,1]), self.twist_degree_range * self.joint_position_threshold * self.joint_direction)
        self.sub_stages.append([delta_twist_pose, None, np.eye(4), 'local_gripper'])

    def axis_angle_to_matrix(self, axis, angle):

        axis = axis / np.linalg.norm(axis)
        K = np.array([[0, -axis[2], axis[1]],
                      [axis[2], 0, -axis[0]],
                      [-axis[1], axis[0], 0]])
        R = np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * (K @ K)
        T = np.eye(4)
        T[:3, :3] = R
        return T

    def check_completion(self, objects):
        if self.__len__() == 0:
            return True

        succ = False
        if self.step_id >= 0:
            succ = True

        if succ:
            self.step_id += 1

        return succ