import numpy as np

from data_gen_dependencies.action.base import StageTemplate, simple_check_completion, pose_difference


class PickStage(StageTemplate):
    def __init__(self, active_obj_id, passive_obj_id, active_element, passive_element, grasp_pose, extra_params=None, active_obj=None, passive_obj=None, **kwargs):
        super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
        self.grasp_pose_canonical = grasp_pose
        self.use_pre_grasp = False

        self.extra_params = {} if extra_params is None else extra_params
        self.pick_up_step = 999
        self.generate_substage(grasp_pose)
        

    def generate_substage(self, grasp_pose):
        pick_up_distance = self.extra_params.get('pick_up_distance', 0.12)
        pick_up_type = self.extra_params.get('pick_up_type', 'Simple')
        pick_up_direction = self.extra_params.get('pick_up_direction', 'z')
        num_against = self.extra_params.get('against', 0)
        if self.use_pre_grasp:
            pre_grasp_distance = self.extra_params.get('pre_grasp_distance', 0.08)
            # sub-stage-0   moveTo pregrasp pose
            pre_pose = np.array([
                [1., 0, 0, 0],
                [0, 1., 0, 0], 
                [0, 0, 1., -pre_grasp_distance],
                [0, 0, 0, 1]])
            pre_grasp_pose = grasp_pose @ pre_pose   
            self.sub_stages.append([pre_grasp_pose, None, np.eye(4), 'AvoidObs'])
            # sub-stage-1   moveTo grasp pose   
            self.sub_stages.append([grasp_pose, 'close', np.eye(4), 'Simple'])
            self.pick_up_step = 2
        elif num_against > 0:
            for index in range(num_against):
                self.sub_stages.append([grasp_pose, 'open', np.eye(4), 'AvoidObs'])
            self.sub_stages.append([grasp_pose, 'close', np.eye(4), 'AvoidObs'])
            self.pick_up_step = num_against + 1 
            gripper_action = None
            motion_type = pick_up_type
            transform_up = np.eye(4)
            if pick_up_direction == "x":
                transform_up[0,3] = pick_up_distance 
                transform_up[2,3] = 0.02
            elif pick_up_direction == "y":
                transform_up[1,3] = pick_up_distance
            else:
                transform_up[2,3] = pick_up_distance
            self.sub_stages.append([grasp_pose, gripper_action, transform_up, motion_type])
        else:
            # sub-stage-0   moveTo grasp pose
            self.sub_stages.append([grasp_pose, 'close', np.eye(4), 'AvoidObs'])
            self.pick_up_step = 1

            # pick-up
            gripper_action = None
            motion_type = pick_up_type
            transform_up = np.eye(4)
            if pick_up_direction == "x":
                transform_up[0,3] = pick_up_distance 
                transform_up[2,3] = 0.02
            elif pick_up_direction == "y":
                transform_up[1,3] = pick_up_distance
            else:
                transform_up[2,3] = pick_up_distance
            self.sub_stages.append([grasp_pose, gripper_action, transform_up, motion_type])



    def check_completion(self, objects):
        if self.__len__()==0:
            return True
        
        goal_datapack = [self.active_obj_id, self.passive_obj_id] + self.sub_stages[self.step_id]

        succ = True
        if self.step_id < self.pick_up_step:
            succ = simple_check_completion(goal_datapack, objects, is_grasped=True)

        if succ:
            self.step_id += 1
        return succ
    







class GraspStage(PickStage):
    def generate_substage(self, grasp_pose):
        motion_type = self.extra_params.get('move_type', 'AvoidObs')

        if self.extra_params.get('use_pre_grasp', False):
            pre_pose = np.array([
                [1., 0, 0, 0],
                [0, 1., 0, 0], 
                [0, 0, 1., -0.08],
                [0, 0, 0, 1]])
            pre_grasp_pose = grasp_pose @ pre_pose   
            gripper_action = None
            
            self.sub_stages.append([pre_grasp_pose, None, np.eye(4), motion_type])
        # sub-stage-1   moveTo grasp pose

        self.sub_stages.append([grasp_pose, 'close', np.eye(4), motion_type])

        



    def check_completion(self, objects):
        if self.__len__()==0:
            return True
        
        goal_datapack = [self.active_obj_id, self.passive_obj_id] + self.sub_stages[self.step_id]

        pos_threshold = self.extra_params.get('position_threshold', 0.06)
        succ = simple_check_completion(goal_datapack, objects, pos_threshold=pos_threshold, is_grasped=True)

        if self.step_id >= 2:
            gripper_pose = objects['gripper'].obj_pose
            grasp_pose = objects[self.passive_obj_id].obj_pose @ self.grasp_pose_canonical
            pos_diff, _ = pose_difference(gripper_pose, grasp_pose)
            grasp_succ = pos_diff < pos_threshold
            succ = succ and grasp_succ

        if succ:
            self.step_id += 1
        return succ
    

class HookStage(PickStage):
    def generate_substage(self, grasp_pose, *args, **kwargs):
        self.sub_stages.append([grasp_pose, None, np.eye(4), 'AvoidObs'])