#
# 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 pytest
import torch

# CuRobo
from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel, CudaRobotModelConfig
from curobo.cuda_robot_model.urdf_kinematics_parser import UrdfKinematicsParser
from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose
from curobo.util_file import (
    file_exists,
    get_assets_path,
    get_robot_configs_path,
    join_path,
    load_yaml,
)

try:
    # CuRobo
    from curobo.cuda_robot_model.usd_kinematics_parser import UsdKinematicsParser
except ImportError:
    pytest.skip("usd-core is not available, skipping USD tests", allow_module_level=True)


def check_usd_file_exists():
    robot_file = "franka.yml"
    robot_params = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
    robot_usd = join_path(get_assets_path(), robot_params["kinematics"]["usd_path"])
    return file_exists(robot_usd)


if not check_usd_file_exists():
    pytest.skip("Franka Panda USD is not available, skipping USD tests", allow_module_level=True)


@pytest.fixture(scope="module")
def robot_params_all():
    robot_file = "franka.yml"
    robot_params = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
    robot_params["kinematics"]["ee_link"] = "panda_link7"
    robot_params["kinematics"]["base_link"] = "panda_link0"
    robot_params["kinematics"]["collision_link_names"] = []
    robot_params["kinematics"]["lock_joints"] = None
    robot_params["kinematics"]["extra_links"] = None

    return robot_params


@pytest.fixture(scope="module")
def usd_parser(robot_params_all):
    robot_params = robot_params_all["kinematics"]
    robot_usd = join_path(get_assets_path(), robot_params["usd_path"])
    kinematics_parser = UsdKinematicsParser(
        usd_path=robot_usd,
        flip_joints=robot_params["usd_flip_joints"],
        usd_robot_root=robot_params["usd_robot_root"],
    )
    return kinematics_parser


@pytest.fixture(scope="module")
def urdf_parser(robot_params_all):
    robot_params = robot_params_all["kinematics"]
    robot_urdf = join_path(get_assets_path(), robot_params["urdf_path"])
    kinematics_parser = UrdfKinematicsParser(robot_urdf)
    return kinematics_parser


@pytest.fixture(scope="module")
def retract_state(robot_params_all):
    tensor_args = TensorDeviceType()
    q = tensor_args.to_device(robot_params_all["kinematics"]["cspace"]["retract_config"]).view(
        1, -1
    )
    return q


@pytest.fixture(scope="module")
def usd_cuda_robot(robot_params_all):
    robot_params = robot_params_all["kinematics"]
    robot_params["use_usd_kinematics"] = True

    usd_cuda_config = CudaRobotModelConfig.from_data_dict(robot_params)

    usd_cuda_robot = CudaRobotModel(usd_cuda_config)
    return usd_cuda_robot


@pytest.fixture(scope="module")
def urdf_cuda_robot(robot_params_all):
    robot_params = robot_params_all["kinematics"]
    robot_params["use_usd_kinematics"] = False
    urdf_cuda_config = CudaRobotModelConfig.from_data_dict(robot_params)

    urdf_cuda_robot = CudaRobotModel(urdf_cuda_config)
    return urdf_cuda_robot


def test_chain_parse(urdf_parser, usd_parser, robot_params_all):
    robot_params = robot_params_all

    urdf_chain = urdf_parser.get_chain(
        robot_params["kinematics"]["base_link"], robot_params["kinematics"]["ee_link"]
    )

    usd_chain = usd_parser.get_chain(
        robot_params["kinematics"]["base_link"], robot_params["kinematics"]["ee_link"]
    )
    assert usd_chain == urdf_chain


def test_joint_transform_parse(usd_cuda_robot, urdf_cuda_robot):
    usd_pose = usd_cuda_robot.get_all_link_transforms()
    urdf_pose = urdf_cuda_robot.get_all_link_transforms()
    p, q = usd_pose.distance(urdf_pose)
    p = torch.linalg.norm(q)
    q = torch.linalg.norm(q)
    assert p < 1e-5 and q < 1e-5


def test_basic_ee_pose(usd_cuda_robot, urdf_cuda_robot, retract_state):
    q = retract_state[:, : usd_cuda_robot.get_dof()]

    usd_state = usd_cuda_robot.get_state(q)
    usd_pose = Pose(usd_state.ee_position, usd_state.ee_quaternion)

    urdf_state = urdf_cuda_robot.get_state(q)
    urdf_pose = Pose(urdf_state.ee_position, urdf_state.ee_quaternion)
    p_d, q_d = usd_pose.distance(urdf_pose)
    assert p_d < 1e-5
    assert q_d < 1e-5