#
# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
#
# Third Party
import torch

# CuRobo
from curobo.util.logger import log_warn

try:
    # CuRobo
    from curobo.curobolib import geom_cu

except ImportError:
    log_warn("geom_cu binary not found, jit compiling...")
    # Third Party
    from torch.utils.cpp_extension import load

    # CuRobo
    from curobo.util_file import add_cpp_path

    geom_cu = load(
        name="geom_cu",
        sources=add_cpp_path(
            [
                "geom_cuda.cpp",
                "sphere_obb_kernel.cu",
                "pose_distance_kernel.cu",
                "self_collision_kernel.cu",
            ]
        ),
    )


def get_self_collision_distance(
    out_distance,
    out_vec,
    sparse_index,
    robot_spheres,
    collision_offset,
    weight,
    coll_matrix,
    thread_locations,
    thread_size,
    b_size,
    nspheres,
    compute_grad,
    checks_per_thread=32,
    experimental_kernel=True,
):
    r = geom_cu.self_collision_distance(
        out_distance,
        out_vec,
        sparse_index,
        robot_spheres,
        collision_offset,
        weight,
        coll_matrix,
        thread_locations,
        thread_size,
        b_size,
        nspheres,
        compute_grad,
        checks_per_thread,
        experimental_kernel,
    )

    out_distance = r[0]
    out_vec = r[1]
    return out_distance, out_vec


class SelfCollisionDistance(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        out_distance,
        out_vec,
        sparse_idx,
        robot_spheres,
        sphere_offset,
        weight,
        coll_matrix,
        thread_locations,
        max_thread,
        checks_per_thread: int,
        experimental_kernel: bool,
        return_loss: bool = False,
    ):
        # get batch size
        b, h, n_spheres, _ = robot_spheres.shape
        out_distance, out_vec = get_self_collision_distance(
            out_distance,
            out_vec,
            sparse_idx,
            robot_spheres,  # .view(-1, 4),
            sphere_offset,
            weight,
            coll_matrix,
            thread_locations,
            max_thread,
            b * h,
            n_spheres,
            robot_spheres.requires_grad,
            checks_per_thread,
            experimental_kernel,
        )
        ctx.return_loss = return_loss
        ctx.save_for_backward(out_vec)
        return out_distance

    @staticmethod
    def backward(ctx, grad_out_distance):
        sphere_grad = None
        if ctx.needs_input_grad[3]:
            (g_vec,) = ctx.saved_tensors
            if ctx.return_loss:
                g_vec = g_vec * grad_out_distance.unsqueeze(1)
            sphere_grad = g_vec
        return None, None, None, sphere_grad, None, None, None, None, None, None, None, None


class SelfCollisionDistanceLoss(SelfCollisionDistance):
    @staticmethod
    def backward(ctx, grad_out_distance):
        sphere_grad = None
        if ctx.needs_input_grad[3]:
            (g_vec,) = ctx.saved_tensors
            sphere_grad = g_vec * grad_out_distance.unsqueeze(1)
        return None, None, None, sphere_grad, None, None, None, None, None, None, None


def get_pose_distance(
    out_distance,
    out_position_distance,
    out_rotation_distance,
    out_p_vec,
    out_q_vec,
    out_idx,
    current_position,
    goal_position,
    current_quat,
    goal_quat,
    vec_weight,
    weight,
    vec_convergence,
    run_weight,
    run_vec_weight,
    batch_pose_idx,
    batch_size,
    horizon,
    mode=1,
    num_goals=1,
    write_grad=False,
    write_distance=False,
    use_metric=False,
):
    if batch_pose_idx.shape[0] != batch_size:
        raise ValueError("Index buffer size is different from batch size")

    r = geom_cu.pose_distance(
        out_distance,
        out_position_distance,
        out_rotation_distance,
        out_p_vec,
        out_q_vec,
        out_idx,
        current_position,
        goal_position.view(-1),
        current_quat,
        goal_quat.view(-1),
        vec_weight,
        weight,
        vec_convergence,
        run_weight,
        run_vec_weight,
        batch_pose_idx,
        batch_size,
        horizon,
        mode,
        num_goals,
        write_grad,
        write_distance,
        use_metric,
    )

    out_distance = r[0]
    out_position_distance = r[1]
    out_rotation_distance = r[2]

    out_p_vec = r[3]
    out_q_vec = r[4]

    out_idx = r[5]
    return out_distance, out_position_distance, out_rotation_distance, out_p_vec, out_q_vec, out_idx


def get_pose_distance_backward(
    out_grad_p,
    out_grad_q,
    grad_distance,
    grad_p_distance,
    grad_q_distance,
    pose_weight,
    grad_p_vec,
    grad_q_vec,
    batch_size,
    use_distance=False,
):
    r = geom_cu.pose_distance_backward(
        out_grad_p,
        out_grad_q,
        grad_distance,
        grad_p_distance,
        grad_q_distance,
        pose_weight,
        grad_p_vec,
        grad_q_vec,
        batch_size,
        use_distance,
    )
    return r[0], r[1]


class SdfSphereOBB(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        query_sphere,
        out_buffer,
        grad_out_buffer,
        sparsity_idx,
        weight,
        activation_distance,
        box_accel,
        box_dims,
        box_pose,
        box_enable,
        n_env_obb,
        env_query_idx,
        max_nobs,
        batch_size,
        horizon,
        n_spheres,
        transform_back,
        compute_distance,
        use_batch_env,
        return_loss: bool = False,
    ):
        r = geom_cu.closest_point(
            query_sphere,
            out_buffer,
            grad_out_buffer,
            sparsity_idx,
            weight,
            activation_distance,
            box_accel,
            box_dims,
            box_pose,
            box_enable,
            n_env_obb,
            env_query_idx,
            max_nobs,
            batch_size,
            horizon,
            n_spheres,
            transform_back,
            compute_distance,
            use_batch_env,
        )
        # r[1][r[1]!=r[1]] = 0.0
        ctx.return_loss = return_loss
        ctx.save_for_backward(r[1])
        return r[0]

    @staticmethod
    def backward(ctx, grad_output):
        grad_pt = None
        if ctx.needs_input_grad[0]:
            (r,) = ctx.saved_tensors
            if ctx.return_loss:
                r = r * grad_output.unsqueeze(-1)
            grad_pt = r
        return (
            grad_pt,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
        )


class SdfSweptSphereOBB(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        query_sphere,
        out_buffer,
        grad_out_buffer,
        sparsity_idx,
        weight,
        activation_distance,
        speed_dt,
        box_accel,
        box_dims,
        box_pose,
        box_enable,
        n_env_obb,
        env_query_idx,
        max_nobs,
        batch_size,
        horizon,
        n_spheres,
        sweep_steps,
        enable_speed_metric,
        transform_back,
        compute_distance,
        use_batch_env,
        return_loss: bool = False,
    ):
        r = geom_cu.swept_closest_point(
            query_sphere,
            out_buffer,
            grad_out_buffer,
            sparsity_idx,
            weight,
            activation_distance,
            speed_dt,
            box_accel,
            box_dims,
            box_pose,
            box_enable,
            n_env_obb,
            env_query_idx,
            max_nobs,
            batch_size,
            horizon,
            n_spheres,
            sweep_steps,
            enable_speed_metric,
            transform_back,
            compute_distance,
            use_batch_env,
        )
        ctx.return_loss = return_loss
        ctx.save_for_backward(
            r[1],
        )
        return r[0]

    @staticmethod
    def backward(ctx, grad_output):
        grad_pt = None
        if ctx.needs_input_grad[0]:
            (r,) = ctx.saved_tensors
            if ctx.return_loss:
                r = r * grad_output.unsqueeze(-1)
            grad_pt = r
        return (
            grad_pt,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
            None,
        )
