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

from typing import Optional
import os
import math
import numpy as np
import torch
from omni.isaac.core.robots.robot import Robot
from omni.isaac.core.utils.nucleus import get_assets_root_path
from omni.isaac.core.utils.stage import add_reference_to_stage
from omni.isaac.sensor import Camera
from omniisaacgymenvs.tasks.utils.usd_utils import set_drive

from omni.isaac.core.utils.prims import get_prim_at_path
from pxr import PhysxSchema

class UR5eSweeper(Robot):
    def __init__(
        self,
        prim_path: str,
        name: Optional[str] = "ur5e_sweeper",
        usd_path: Optional[str] = None,
        translation: Optional[torch.tensor] = None,
        orientation: Optional[torch.tensor] = None,
        # rgb_cam_prim_name: Optional[str] = None,    ### BSH
        # depth_cam_prim_name: Optional[str] = None,  ### BSH
    ) -> None:
        """[summary]
        """

        self._usd_path = usd_path
        self._name = name

        self._position = torch.tensor([0, 0.0, 0.0]) if translation is None else translation
        # self._orientation = torch.tensor([0.0, 0.0, 0.0, 1.0]) if orientation is None else orientation
        self._orientation = torch.tensor([1.0, 0.0, 0.0, 0.0]) if orientation is None else orientation

        if self._usd_path is None:
            isaac_root_path = "/home/bak/.local/share/ov/pkg/isaac_sim-2022.2.0/"
            if isaac_root_path is None:
                carb.log_error("Could not find Isaac Sim assets folder")
            # self._usd_path = isaac_root_path + "standalone_examples/user_example/ur5e.usd"
            # self._usd_path = isaac_root_path + "standalone_examples/user_example/ur5e_handeye_gripper.usd"
            # self._usd_path = isaac_root_path + "OmniIsaacGymEnvs/omniisaacgymenvs/skrl_examples/reaching_ur5e/ur5e_tool.usd"
            # self._usd_path = isaac_root_path + "OmniIsaacGymEnvs/omniisaacgymenvs/skrl_examples/moving_target/ur5e_tool_cam.usd"
            self._usd_path = isaac_root_path + "OmniIsaacGymEnvs/omniisaacgymenvs/skrl_examples/sweeping/ur5e/usd/ur5e_sweeper.usd"
            # self._usd_path = isaac_root_path + "OmniIsaacGymEnvs/omniisaacgymenvs/skrl_examples/moving_target/ur5e/usd/ur5e_peg_merged.usd"

        # if rgb_cam_prim_name is None:   ### BSH
        #     self._rgb_cam_prim_path = prim_path + "/realsense/rgb"
        # if depth_cam_prim_name is None:   ### BSH
        #     self._depth_cam_prim_path = prim_path + "/realsense/depth"

        add_reference_to_stage(self._usd_path, prim_path)
        
        super().__init__(
            prim_path=prim_path,
            name=name,
            translation=self._position,
            orientation=self._orientation,
            articulation_controller=None,
        )

        # dof_paths = [
        #     "base_link_inertia/shoulder_pan_joint",
        #     "shoulder_link/shoulder_lift_joint",
        #     "upper_arm_link/elbow_joint",
        #     "forearm_link/wrist_1_joint",
        #     "wrist_1_link/wrist_2_joint",
        #     "wrist_2_link/wrist_3_joint",
        # ]
        dof_paths = [
            "base_link/shoulder_pan_joint",
            "shoulder_link/shoulder_lift_joint",
            "upper_arm_link/elbow_joint",
            "forearm_link/wrist_1_joint",
            "wrist_1_link/wrist_2_joint",
            "wrist_2_link/wrist_3_joint",
        ]

        drive_type = ["angular"] * 6
        # default_dof_pos = [math.degrees(x) for x in [0.0, -1.75, 1.05, -1.57, -1.57, 0.0]]
        default_dof_pos = [math.degrees(x) for x in [0.0, -1.75, 1.05, -1.57, -1.57, 1.57]]
        # default_dof_pos = [math.degrees(x) for x in [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]
        # default_dof_pos [0.0, -100.0, 60.0, -90.0, -90.0, 90.0]
        # default_dof_pos = [math.degrees(x) for x in [0.0, -60.0, -60.0, -90.0, 90.0, 0.0]]
        stiffness = [400*np.pi/180] * 6
        damping = [80*np.pi/180] * 6
        max_force = [87, 87, 87, 12, 12, 12]
        max_velocity = [math.degrees(x) for x in [2.175, 2.175, 2.175, 2.61, 2.61, 2.61]]

        # self._rgb_cam = Camera(prim_path=self._rgb_cam_prim_path,
        #                        frequency=30, resolution=(1920, 1080)) ### BSH
        # self._depth_cam = Camera(prim_path=self._depth_cam_prim_path,
        #                         #  frequency=30, resolution=(1920, 1080)) ### BSH
        #                         frequency=30, resolution=(1280, 720)) ### BSH

        for i, dof in enumerate(dof_paths):
            set_drive(
                prim_path=f"{self.prim_path}/{dof}",
                drive_type=drive_type[i],
                target_type="position",
                target_value=default_dof_pos[i],
                stiffness=stiffness[i],
                damping=damping[i],
                max_force=max_force[i]
            )

            PhysxSchema.PhysxJointAPI(get_prim_at_path(f"{self.prim_path}/{dof}")).CreateMaxJointVelocityAttr().Set(max_velocity[i])