import numpy as np

from task_gen_dependencies.sdf import get_distance_with_sdf
from task_gen_dependencies.transform_utils import transform_points, random_point


def cost_without_collision(obj_active, obj_passive, threshold=1, **kwargs):
    collision_points = obj_active.collision_points
    active2passive = np.linalg.inv(obj_passive.obj_pose) @ obj_active.obj_pose
    sdf_distance = get_distance_with_sdf(collision_points, active2passive, obj_passive.sdf)

    cost = np.sum(sdf_distance < 0)
    print('sdf:', cost, sdf_distance.min(), sdf_distance.max())
    return cost * 10


def cost_A_out_B(obj_active, obj_passive, safe_distance, **kwargs):
    activate_point = obj_active.obj_pose[:3, 3]
    passive_point = obj_passive.obj_pose[:3, 3]
    distance_vector = activate_point[:2] - passive_point[:2]
    distance = np.linalg.norm(distance_vector)

    cost = 0
    if distance < safe_distance:
        cost += 10
    return cost


def cost_A_on_B(obj_active, obj_passive, **kwargs):
    active_pts = transform_points(obj_active.anchor_points['buttom'], obj_active.obj_pose)[0]
    passive_pts = transform_points(obj_passive.anchor_points['top'], obj_passive.obj_pose)[0]

    # active_pts = transform_points(obj_active.anchor_points['buttom'], obj_active.obj_pose)[0]
    # passive_pts = transform_points(obj_passive.anchor_points['top'], obj_passive.obj_pose)[0]

    cost_add = 0
    # if activate_point[2] <= passive_point[2]:
    #     cost_add += 10  # Large penalty if the constraint is violated
    distance_vector = active_pts - passive_pts
    # import ipdb; ipdb.set_trace()
    distance_cost = np.linalg.norm(distance_vector)
    return cost_add + distance_cost


def cost_A_in_B(opt_pose_active_homo, opt_pose_passive_homo, selected_points_activate_sam, selected_points_passive_sam,
                selected_points_passive_bottom):
    # import pdb;pdb.set_trace()

    transformed_points_activate = batch_transform_points(selected_points_activate_sam, opt_pose_active_homo)
    transformed_points_activate = transformed_points_activate.reshape(-1, 3)

    transformed_points_passive = selected_points_passive_sam
    transformed_points_passive = transformed_points_passive.reshape(-1, 3)

    transformed_points_passive_bot = selected_points_passive_bottom
    transformed_points_passive_bot = transformed_points_passive_bot.reshape(-1, 3)
    # Get the lowest point of the active object
    activate_point = random_point(transformed_points_activate, 3)
    # Get the highest point of the passive object
    passive_point_top = random_point(transformed_points_passive, 3)

    passive_point_bot = random_point(transformed_points_passive_bot, 10)  # 约束靠近中心，避免求解出来碰撞
    cost_add = 0
    if activate_point[2] >= passive_point_top[2]:
        cost_add += 10  # Large penalty if the constraint is violated
    distance_vector = activate_point - passive_point_bot
    distance_cost = np.linalg.norm(distance_vector)
    return cost_add + distance_cost


CONSTRAINT_COST = {
    'wo_collision': cost_without_collision,
    "out": cost_A_out_B,
    "on": cost_A_on_B,
    "in": cost_A_in_B,
    "faceto": NotImplemented,
    'backto': NotImplemented
}


def get_cost_func(constraint):
    if constraint not in constraint:
        raise NotImplementedError("Constraint {} not implemented".format(constraint))
    return CONSTRAINT_COST[constraint]