# -*- coding: utf-8 -*-
import copy
import json
import os
from sqlite3 import dbapi2
import time

from pyboot.utils.log import Log
import trimesh

curPath = os.path.abspath(os.path.dirname(__file__))
rootPath = os.path.split(curPath)[0]
import numpy as np

from data_gen_dependencies.transforms import calculate_rotation_matrix, rotate_around_axis
from data_gen_dependencies.object import OmniObject, transform_coordinates_3d
from data_gen_dependencies.fix_rotation import rotate_along_axis

from data_gen_dependencies.action import build_stage

from scipy.interpolate import interp1d
from scipy.spatial.transform import Rotation as R


def draw_axis(robot, pose):
    if pose.shape[0] == 4 and pose.shape[1] == 4:
        pose = np.concatenate([pose[:3, 3], R.from_matrix(pose[:3, :3]).as_euler('xyz')])

    def generate_orthogonal_lines(point, angles, scale):
        assert len(point) == 3 and len(angles) == 3

        # Define rotation matrices for each axis
        Rx = np.array([[1, 0, 0],
                       [0, np.cos(angles[0]), -np.sin(angles[0])],
                       [0, np.sin(angles[0]), np.cos(angles[0])]])

        Ry = np.array([[np.cos(angles[1]), 0, np.sin(angles[1])],
                       [0, 1, 0],
                       [-np.sin(angles[1]), 0, np.cos(angles[1])]])

        Rz = np.array([[np.cos(angles[2]), -np.sin(angles[2]), 0],
                       [np.sin(angles[2]), np.cos(angles[2]), 0],
                       [0, 0, 1]])

        # Combined rotation matrix
        R = Rz @ Ry @ Rx

        # Define the unit vectors for the axes
        unit_vectors = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])

        # Rotate the unit vectors
        rotated_vectors = R @ unit_vectors.T

        # Draw the lines
        lines = []
        for vec in rotated_vectors.T:
            start_point = point
            end_point = point + vec * scale
            lines.append(np.concatenate([start_point, end_point]))
        return np.stack(lines)

    draw_colors = [
        [0, 0, 1],
        [0, 1, 0],
        [1, 0, 0],
    ]
    line_sizes = [1, 1, 1]
    xyz_axis = generate_orthogonal_lines(point=pose[:3], angles=pose[3:], scale=0.1)
    start_point = xyz_axis[:, :3]
    end_point = xyz_axis[:, 3:]
    robot.client.DrawLine(start_point, end_point, draw_colors, line_sizes)


def load_gripper_width_interp_func(file):
    data = json.load(open(file, 'r'))

    digits = []
    widths = []
    depths = []
    for frame in data:
        digits.append(frame["angel"])
        widths.append(frame["width"])
        depths.append(frame["depth"])
    digits = np.array(digits)
    widths = np.array(widths)
    depths = np.array(depths)

    func_width2depth = interp1d(widths, depths, kind="linear", fill_value="extrapolate")
    return func_width2depth


func_width2depth = load_gripper_width_interp_func(os.path.dirname(__file__) + "/gripper_width_120s_fixed.json")


def format_object(obj, distance, type='active'):
    if obj is None:
        return None
    xyz, direction = obj.xyz, obj.direction

    direction = direction / np.linalg.norm(direction) * distance
    type = type.lower()
    if type == 'active':
        xyz_start = xyz
        xyz_end = xyz_start + direction
    elif type == 'passive' or type == 'plane':
        xyz_end = xyz
        xyz_start = xyz_end - direction

    part2obj = np.eye(4)
    part2obj[:3, 3] = xyz_start
    obj.obj2part = np.linalg.inv(part2obj)

    obj_info = {
        'pose': obj.obj_pose,
        'length': obj.obj_length,
        'xyz_start': xyz_start,
        'xyz_end': xyz_end,
        'obj2part': obj.obj2part
    }
    return obj_info


def obj2world(obj_info):
    obj_pose = obj_info['pose']
    obj_length = obj_info['length']
    obj2part = obj_info['obj2part']
    xyz_start = obj_info['xyz_start']
    xyz_end = obj_info['xyz_end']

    arrow_in_obj = np.array([xyz_start, xyz_end]).transpose(1, 0)
    arrow_in_world = transform_coordinates_3d(arrow_in_obj, obj_pose).transpose(1, 0)

    xyz_start_world, xyz_end_world = arrow_in_world
    direction_world = xyz_end_world - xyz_start_world
    direction_world = direction_world / np.linalg.norm(direction_world)

    obj_info_world = {
        'pose': obj_pose,
        'length': obj_length,
        'obj2part': obj2part,
        'xyz_start': xyz_start_world,
        'xyz_end': xyz_end_world,
        'direction': direction_world,

    }
    return obj_info_world


def get_aligned_fix_pose(active_obj, passive_obj, distance=0.01, N=1):
    try:
        active_object = format_object(active_obj, type='active', distance=distance)
        passive_object = format_object(passive_obj, type='passive', distance=distance)
    except:
        print('error')
    active_obj_world = obj2world(active_object)
    current_obj_pose = active_obj_world['pose']
    if passive_object is None:
        return current_obj_pose[np.newaxis, ...]

    passive_obj_world = obj2world(passive_object)
    passive_obj_world['direction'] = passive_obj.direction

    R = calculate_rotation_matrix(active_obj_world['direction'], passive_obj_world['direction'])
    T = passive_obj_world['xyz_end'] - R @ active_obj_world['xyz_start']
    transform_matrix = np.eye(4)
    transform_matrix[:3, :3] = R
    transform_matrix[:3, 3] = T
    target_obj_pose = transform_matrix @ current_obj_pose
    target_obj_pose[:3, 3] = passive_obj.obj_pose[:3, 3]

    poses = []
    for angle in [i * 360 / N for i in range(N)]:
        pose_rotated = rotate_around_axis(
            target_obj_pose,
            passive_obj_world['xyz_start'],
            passive_obj_world['direction'],
            angle)
        poses.append(pose_rotated)
    return np.stack(poses)


def get_aligned_pose(active_obj, passive_obj, distance=0.01, N=1):
    try:
        active_object = format_object(active_obj, type='active', distance=distance)
        passive_object = format_object(passive_obj, type='passive', distance=distance)
    except:
        print('error')

    active_obj_world = obj2world(active_object)
    current_obj_pose = active_obj_world['pose']
    if passive_object is None:
        return current_obj_pose[np.newaxis, ...]

    passive_obj_world = obj2world(passive_object)

    R = calculate_rotation_matrix(active_obj_world['direction'], passive_obj_world['direction'])
    T = passive_obj_world['xyz_end'] - R @ active_obj_world['xyz_start']
    transform_matrix = np.eye(4)
    transform_matrix[:3, :3] = R
    transform_matrix[:3, 3] = T
    target_obj_pose = transform_matrix @ current_obj_pose

    poses = []
    for angle in [i * 360 / N for i in range(N)]:
        pose_rotated = rotate_around_axis(
            target_obj_pose,
            passive_obj_world['xyz_start'],
            passive_obj_world['direction'],
            angle)
        poses.append(pose_rotated)
    return np.stack(poses)


def load_task_solution(task_info):
    Log.info('Loading task solution')
    # task_info = json.load(open('%s/task_info.json'%task_dir, 'rb'))
    stages = task_info['stages']

    objects = {
        'gripper': OmniObject('gripper')
    }

    for i, obj_info in enumerate(task_info['objects']):
        Log.info(f'Loading object: {obj_info["object_id"]} [Loading task solution: {i+1}/{len(task_info["objects"])}]')
        obj_id = obj_info['object_id']
        if obj_id == 'fix_pose':
            obj = OmniObject('fix_pose')
            if 'position' not in obj_info or 'direction' not in obj_info:
                print(f"Error: Missing position/direction in object {obj_id}")
                continue
            obj.set_pose(np.array(obj_info['position']), np.array([0.001, 0.001, 0.001]))
            obj.elements = {
                'active': {},
                'passive': {
                    'place': {
                        'default': [
                            {
                                'xyz': np.array([0, 0, 0]),
                                'direction': np.array(obj_info['direction'])
                            }
                        ]
                    }
                }
            }
            objects[obj_id] = obj
        else:
            obj_dir = obj_info["data_info_dir"]
            obj = OmniObject.from_obj_dir(obj_dir, obj_info=obj_info)
            objects[obj_id] = obj

            if hasattr(obj, 'part_ids'):
                for part_id in obj.part_ids:
                    id = obj_id + '/%s' % part_id
                    objects[id] = copy.deepcopy(obj)
                    objects[id].name = id
                    #====== DEBUG START ======
                    print(f'debug: ')
                    import ipdb;ipdb.set_trace()
                    #======  DEBUG END  ======
                    if hasattr(obj, 'part_joint_limits') and obj.part_joint_limits is not None:
                        objects[id].part_joint_limit = obj.part_joint_limits[part_id]
                if len(obj.part_ids):
                    del objects[obj_id]
    Log.success('Finished Loading task solution')
    return stages, objects


def parse_stage(stage, objects):
    action = stage['action']
    if action in ['reset']:
        return action, 'gripper', 'gripper', None, None, None, None
    active_obj_id = stage['active']['object_id']
    if 'part_id' in stage['active']:
        active_obj_id += '/%s' % stage['active']['part_id']

    passive_obj_id = stage['passive']['object_id']
    if 'part_id' in stage['passive']:
        passive_obj_id += '/%s' % stage['passive']['part_id']

    active_obj = objects[active_obj_id]
    passive_obj = objects[passive_obj_id]

    single_obj = action in ['pull', 'rotate', 'slide', 'shave', 'brush', 'wipe']

    def _load_element(obj, type):
        if action in ['pick', 'hook']:
            action_mapped = 'grasp'
        else:
            action_mapped = action
        if action_mapped == 'grasp' and type == 'active':
            return None, None
        elif obj.name == 'gripper':
            element = obj.elements[type][action_mapped]
            return element, 'default'
        primitive = stage[type]['primitive'] if stage[type]['primitive'] is not None else 'default'
        if primitive != 'default' or (action_mapped == 'grasp' and type == 'passive'):
            if action_mapped not in obj.elements[type]:
                print('No %s element for %s' % (action_mapped, obj.name))
                return None, None
            element = obj.elements[type][action_mapped][primitive]
        else:
            element = []
            for primitive in obj.elements[type][action_mapped]:
                _element = obj.elements[type][action_mapped][primitive]
                if isinstance(_element, list):
                    element += _element
                else:
                    element.append(_element)
        return element, primitive

    passive_element, passive_primitive = _load_element(passive_obj, 'passive')
    if not single_obj:
        active_element, active_primitive = _load_element(active_obj, 'active')
    else:
        active_element, active_primitive = passive_element, passive_primitive
    return action, active_obj_id, passive_obj_id, active_element, passive_element, active_primitive, passive_primitive


def select_obj(objects, stages, robot):
    gripper2obj = None
    extra_params = stages[0].get('extra_params', {})
    arm = extra_params.get('arm', 'right')
    current_gripper_pose = robot.get_ee_pose(id=arm)

    ''' 初筛抓取pose,得到 grasp_poses_canonical, grasp_poses '''
    grasp_stage_id = None

    if stages[0]['action'] in ['push', 'reset', 'click']:
        gripper2obj = current_gripper_pose

    elif stages[0]['action'] in ['pick', 'grasp', 'hook']:
        action = stages[0]['action']

        ''' 筛掉无IK解的grasp pose '''
        grasp_stage_id = 0
        grasp_stage = parse_stage(stages[0], objects)
        _, _, passive_obj_id, _, passive_element, _, _ = grasp_stage
        grasp_obj_id = passive_obj_id
        grasp_poses_canonical = passive_element['grasp_pose'].copy()
        grasp_widths = passive_element['width']

        grasp_poses_canonical[:, :3, :3] = grasp_poses_canonical[:, :3, :3] @ robot.robot_gripper_2_grasp_gripper[
            np.newaxis, ...]

        grasp_poses_canonical_flip = []
        for _i in range(grasp_poses_canonical.shape[0]):
            grasp_poses_canonical_flip.append(rotate_along_axis(grasp_poses_canonical[_i], 180, 'z', use_local=True))
        grasp_poses_canonical_flip = np.stack(grasp_poses_canonical_flip)
        grasp_poses_canonical = np.concatenate([grasp_poses_canonical, grasp_poses_canonical_flip], axis=0)
        grasp_widths = np.concatenate([grasp_widths, grasp_widths], axis=0)
        grasp_poses = objects[passive_obj_id].obj_pose[np.newaxis, ...] @ grasp_poses_canonical

        # filter with IK-checking
        ik_success, _ = robot.solve_ik(grasp_poses, ee_type='gripper', arm=arm, type='Simple')
        grasp_poses_canonical, grasp_poses = grasp_poses_canonical[ik_success], grasp_poses[ik_success]
        grasp_widths = grasp_widths[ik_success]

        if len(grasp_poses) == 0:
            print(action, 'No grasp_gripper_pose can pass Isaac IK')
            return []

    ''' 基于有IK解的grasp pose分数,选择最优的passive primitive element,同时选出最优的一个grasp pose'''
    if grasp_stage_id is not None:
        next_stage_id = grasp_stage_id + 1
        if next_stage_id < len(stages):
            action, active_obj_id, passive_obj_id, active_elements, passive_elements, active_primitive, passive_primitive = parse_stage(
                stages[next_stage_id], objects)

            single_obj = active_obj_id == passive_obj_id

            active_obj = objects[active_obj_id]
            passive_obj = objects[passive_obj_id]
            passive_element = passive_elements[np.random.choice(len(passive_elements))]

            if action == 'place':  # TODO A2D暂时这样搞，Franka要取消
                # import ipdb; ipdb.set_trace()
                obj_pose = active_obj.obj_pose
                mesh = trimesh.load(active_obj.info['mesh_file'], force="mesh")
                mesh.apply_scale(0.001)
                mesh.apply_transform(obj_pose)
                pts, _ = trimesh.sample.sample_surface(mesh, 200)  # 表面采样
                xyz = np.array([np.mean(pts[:, 0]), np.mean(pts[:, 1]), np.percentile(pts[:, 2], 1)])

                direction = np.array([0, 0, -1])
                xyz_canonical = (np.linalg.inv(obj_pose) @ np.array([*xyz, 1]))[:3]
                direction_canonical = (np.linalg.inv(obj_pose) @ np.array([*direction, 0]))[:3]
                active_elements = [{'xyz': xyz_canonical, 'direction': direction_canonical}]

            # import ipdb; ipdb.set_trace()
            t0 = time.time()
            element_ik_score = []
            grasp_pose_ik_score = []
            for active_element in active_elements:
                # interaction between two rigid objects
                obj_pose = active_obj.obj_pose

                N_align = 12
                if not single_obj:
                    active_obj.xyz, active_obj.direction = active_element['xyz'], active_element['direction']
                    if passive_obj_id == 'fix_pose':
                        passive_obj.xyz, passive_obj.direction = passive_element['xyz'], passive_element['direction']
                        target_obj_poses = get_aligned_fix_pose(active_obj, passive_obj, N=N_align)
                    else:
                        passive_obj.xyz, passive_obj.direction = passive_element['xyz'], passive_element['direction']
                        target_obj_poses = get_aligned_pose(active_obj, passive_obj, N=N_align)
                else:  # 物体自身移动
                    transform = np.eye(4)
                    transform[:3, 3] = active_element['xyz']
                    target_obj_poses = (obj_pose @ transform)[np.newaxis, ...]
                    N_align = 1

                N_obj_pose = target_obj_poses.shape[0]
                N_grasp_pose = grasp_poses_canonical.shape[0]
                target_gripper_poses = (
                            target_obj_poses[:, np.newaxis, ...] @ grasp_poses_canonical[np.newaxis, ...]).reshape(-1,
                                                                                                                   4, 4)

                ik_success, _ = robot.solve_ik(target_gripper_poses, ee_type='gripper', type='Simple', arm=arm)
                element_ik_score.append(np.max(ik_success.reshape(N_obj_pose, N_grasp_pose).sum(axis=1)))

                grasp_pose_ik = ik_success.reshape(N_obj_pose, N_grasp_pose)
                grasp_pose_ik_score.append(np.sum(grasp_pose_ik, axis=0))

            # print("ik solve time: ", time.time() - t0)
            best_element_id = np.argmax(element_ik_score)
            best_active_element = active_elements[best_element_id]

            if not single_obj:
                active_obj.elements['active'][action] = {active_primitive: best_active_element}

            grasp_ik_score = grasp_pose_ik_score[best_element_id]

            # import ipdb;ipdb.set_trace()
            _mask = grasp_ik_score >= np.median(grasp_ik_score) / 2
            best_grasp_poses = grasp_poses[_mask]

            # downsample grasp pose
            downsample_num = 100
            grasps_num = best_grasp_poses.shape[0]
            selected_index = np.random.choice(grasps_num, size=min(downsample_num, grasps_num), replace=False)
            best_grasp_poses = best_grasp_poses[selected_index]
            # if best_grasp_poses.shape[0] > downsample_num:
            #     best_grasp_poses = best_grasp_poses[:downsample_num]
            if best_grasp_poses.shape[0] > 1:
                joint_names = robot.joint_names[arm]

                ik_success, ik_info = robot.solve_ik(best_grasp_poses, ee_type='gripper', type='Simple', arm=arm)
                best_grasp_poses = best_grasp_poses[ik_success]
                ik_joint_positions = ik_info["joint_positions"][ik_success]
                ik_joint_names = ik_info["joint_names"][ik_success]
                ik_jacobian_score = ik_info["jacobian_score"][ik_success]

                if len(best_grasp_poses) == 0:
                    print(action, 'No best_grasp_poses can pass curobo IK')
                    return []
                target_joint_positions = []
                for ik_joint_position, ik_joint_name in zip(ik_joint_positions, ik_joint_names):
                    temp_target_joint_positions = []
                    for joint_name in joint_names:
                        temp_target_joint_positions.append(ik_joint_position[list(ik_joint_name).index(joint_name)])
                    target_joint_positions.append(np.array(temp_target_joint_positions))
                target_joint_positions = np.array(target_joint_positions)

                # choose target pose based on the ik jacobian score and ik joint positions
                cur_joint_states = robot.client.get_joint_positions().states
                cur_joint_positions = []
                for key in cur_joint_states:
                    if key.name in joint_names:
                        cur_joint_positions.append(key.position)
                cur_joint_positions = np.array(cur_joint_positions)
                joint_pos_dist = np.linalg.norm(target_joint_positions - cur_joint_positions[np.newaxis, :], axis=1)
                # joint_pos_dist = (joint_pos_dist - np.min(joint_pos_dist))/(np.max(joint_pos_dist) - np.min(joint_pos_dist))
                # ik_jacobian_score = (ik_jacobian_score - np.min(ik_jacobian_score))/(np.max(ik_jacobian_score) - np.min(ik_jacobian_score))
                cost = joint_pos_dist  # - ik_jacobian_score
                idx_sorted = np.argsort(cost)
                best_grasp_pose = best_grasp_poses[idx_sorted][0]

            else:
                best_grasp_pose = best_grasp_poses[0]

            best_grasp_pose_canonical = np.linalg.inv(objects[grasp_obj_id].obj_pose) @ best_grasp_pose
            gripper2obj = best_grasp_pose_canonical
        else:
            gripper2obj = grasp_poses_canonical[0]
    return gripper2obj


def split_grasp_stages(stages):
    split_stages = []
    i = 0
    while i < len(stages):
        if stages[i]['action'] in ['pick', 'grasp', 'hook']:
            if (i + 1) < len(stages) and stages[i + 1]['action'] not in ['pick', 'grasp', 'hook']:
                split_stages.append([stages[i], stages[i + 1]])
                i += 2
            else:
                split_stages.append([stages[i]])
                i += 1
        else:
            split_stages.append([stages[i]])
            i += 1
    return split_stages


def generate_action_stages(objects, all_stages, robot):
    split_stages = split_grasp_stages(all_stages)

    action_stages = []
    for stages in split_stages:
        gripper2obj = select_obj(objects, stages, robot)
        if gripper2obj is None or len(gripper2obj) == 0:
            print('No gripper2obj pose can pass IK')
            gripper2obj = select_obj(objects, stages, robot)
            return []
        for stage in stages:
            extra_params = stage.get('extra_params', {})
            arm = extra_params.get('arm', 'right')
            current_gripper_pose = robot.get_ee_pose(id=arm)
            action, active_obj_id, passive_obj_id, active_elements, passive_elements, active_primitive, passive_primitive = parse_stage(
                stage, objects)
            active_obj = objects[active_obj_id]
            passive_obj = objects[passive_obj_id]

            single_obj = active_obj_id == passive_obj_id

            substages = None
            if action in ['reset']:
                substages = True
            elif action in ['pick', 'grasp', 'hook']:
                substages = build_stage(action)(active_obj_id, passive_obj_id, active_elements, passive_elements,
                                                gripper2obj, extra_params=stage.get('extra_params', None))
            elif action in ['slide', 'shave', 'brush', 'wipe', 'pull', 'rotate', 'pull', 'move',
                            'reset']:  # grasp + 物体自身运动
                passive_element = passive_elements[np.random.choice(len(passive_elements))]
                substages = build_stage(action)(
                    active_obj_id=active_obj_id,
                    passive_obj_id=passive_obj_id,
                    passive_element=passive_element
                )
            else:
                passive_element = passive_elements[np.random.choice(len(passive_elements))]
                # active_element = active_elements[np.random.choice(len(active_elements))] if isinstance(active_elements, list) else active_elements
                if not isinstance(active_elements, list):
                    active_elements = [active_elements]

                for active_element in active_elements:
                    joint_names = robot.joint_names[arm]

                    # interaction between two rigid objects
                    obj_pose = active_obj.obj_pose
                    anchor_pose = passive_obj.obj_pose
                    current_obj_pose_canonical = np.linalg.inv(anchor_pose) @ obj_pose
                    active_obj.xyz, active_obj.direction = active_element['xyz'], active_element['direction']
                    passive_obj.xyz, passive_obj.direction = passive_element['xyz'], passive_element['direction']
                    if active_obj.name == 'gripper':
                        gripper2obj = np.eye(4)

                    if passive_obj_id == 'fix_pose':
                        # import ipdb;ipdb.set_trace()
                        target_obj_poses = get_aligned_fix_pose(active_obj, passive_obj, N=18)
                    else:
                        target_obj_poses = get_aligned_pose(active_obj, passive_obj, N=18)
                    target_gripper_poses = target_obj_poses @ gripper2obj[np.newaxis, ...]
                    # import ipdb;ipdb.set_trace()

                    # downsample target_gripper_poses
                    downsample_num = 100
                    if target_gripper_poses.shape[0] > downsample_num:
                        target_gripper_poses = target_gripper_poses[:downsample_num]
                    ik_success, ik_info = robot.solve_ik(target_gripper_poses, ee_type='gripper', type='AvoidObs',
                                                         arm=arm)
                    target_gripper_poses_pass_ik = target_gripper_poses[ik_success]
                    ik_joint_positions = ik_info["joint_positions"][ik_success]
                    ik_joint_names = ik_info["joint_names"][ik_success]
                    ik_jacobian_score = ik_info["jacobian_score"][ik_success]

                    if len(target_gripper_poses_pass_ik) == 0:
                        print(action, ': No target_obj_pose can pass isaac curobo IK')
                        continue

                    target_joint_positions = []
                    for ik_joint_position, ik_joint_name in zip(ik_joint_positions, ik_joint_names):
                        temp_target_joint_positions = []
                        for joint_name in joint_names:
                            temp_target_joint_positions.append(ik_joint_position[list(ik_joint_name).index(joint_name)])
                        target_joint_positions.append(np.array(temp_target_joint_positions))
                    target_joint_positions = np.array(target_joint_positions)

                    # choose target pose based on the jacobian score and ik joint positions difference
                    cur_joint_states = robot.client.get_joint_positions().states
                    cur_joint_positions = []
                    for key in cur_joint_states:
                        if key.name in joint_names:
                            cur_joint_positions.append(key.position)
                    cur_joint_positions = np.array(cur_joint_positions)
                    joint_pos_dist = np.linalg.norm(target_joint_positions - cur_joint_positions[np.newaxis, :], axis=1)
                    # joint_pos_dist = (joint_pos_dist - np.min(joint_pos_dist))/(np.max(joint_pos_dist) - np.min(joint_pos_dist))
                    # ik_jacobian_score = (ik_jacobian_score - np.min(ik_jacobian_score))/(np.max(ik_jacobian_score) - np.min(ik_jacobian_score))
                    cost = joint_pos_dist  # - ik_jacobian_score
                    idx_sorted = np.argsort(cost)

                    best_target_gripper_pose = target_gripper_poses_pass_ik[idx_sorted][0]

                    best_target_obj_pose = best_target_gripper_pose @ np.linalg.inv(gripper2obj)
                    target_obj_pose_canonical = np.linalg.inv(
                        anchor_pose) @ best_target_obj_pose  # TODO 暂时只取一个可行解，后面要结合grasp pose做整条trajectory的joint solve

                    part2obj = np.eye(4)
                    part2obj[:3, 3] = active_obj.xyz
                    obj2part = np.linalg.inv(part2obj)

                    substages = build_stage(action)(
                        active_obj_id=active_obj_id,
                        passive_obj_id=passive_obj_id,
                        target_pose=target_obj_pose_canonical,
                        current_pose=current_obj_pose_canonical,
                        obj2part=obj2part,
                        vector_direction=passive_element['direction'],
                        passive_element=passive_element,
                        extra_params=stage.get('extra_params', {})
                    )
                    break

            if substages is None:
                print(action, ': No target_obj_pose can pass IK')
                return []
            action_stages.append((action, substages))

    return action_stages


def solve_target_gripper_pose(stage, objects):
    active_obj_ID, passive_obj_ID, target_pose_canonical, gripper_action, transform_world, motion_type = stage

    anchor_pose = objects[passive_obj_ID].obj_pose

    if motion_type == 'Trajectory':
        assert len(target_pose_canonical.shape) == 3, 'The target_pose should be a list of poses'
        target_pose = anchor_pose[np.newaxis, ...] @ target_pose_canonical
        target_pose = transform_world[np.newaxis, ...] @ target_pose
    else:
        target_pose = anchor_pose @ target_pose_canonical
        target_pose = transform_world @ target_pose
    assert 'gripper' in objects, 'The gripper should be the first one in the object list'
    current_gripper_pose = objects['gripper'].obj_pose

    if active_obj_ID == 'gripper':
        target_gripper_pose = target_pose
    else:
        current_obj_pose = objects[active_obj_ID].obj_pose
        gripper2obj = np.linalg.inv(current_obj_pose) @ current_gripper_pose
        if len(target_pose.shape) == 3:
            gripper2obj = gripper2obj[np.newaxis, ...]

        target_obj_pose = target_pose
        target_gripper_pose = target_obj_pose @ gripper2obj

    return target_gripper_pose