#
# 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.
#
""" This example shows how to use cuRobo's kinematics to generate a mask. """


# Standard Library
import time

# Third Party
import imageio
import numpy as np
import torch
import torch.autograd.profiler as profiler
from nvblox_torch.datasets.mesh_dataset import MeshDataset
from torch.profiler import ProfilerActivity, profile, record_function

# CuRobo
from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel
from curobo.geom.types import PointCloud, WorldConfig
from curobo.types.base import TensorDeviceType
from curobo.types.camera import CameraObservation
from curobo.types.math import Pose
from curobo.types.robot import RobotConfig
from curobo.types.state import JointState
from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
from curobo.wrap.model.robot_segmenter import RobotSegmenter

torch.manual_seed(30)

torch.backends.cudnn.benchmark = True

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True


def create_render_dataset(
    robot_file,
    save_debug_data: bool = False,
    fov_deg: float = 60,
    n_frames: int = 20,
    retract_delta: float = 0.0,
):
    # load robot:
    robot_dict = load_yaml(join_path(get_robot_configs_path(), robot_file))
    robot_dict["robot_cfg"]["kinematics"]["load_link_names_with_mesh"] = True
    robot_dict["robot_cfg"]["kinematics"]["load_meshes"] = True

    robot_cfg = RobotConfig.from_dict(robot_dict["robot_cfg"])

    kin_model = CudaRobotModel(robot_cfg.kinematics)

    q = kin_model.retract_config

    q += retract_delta

    meshes = kin_model.get_robot_as_mesh(q)

    world = WorldConfig(mesh=meshes[:])
    world_table = WorldConfig.from_dict(
        load_yaml(join_path(get_world_configs_path(), "collision_test.yml"))
    )
    world_table.cuboid[0].dims = [0.5, 0.5, 0.1]
    world.add_obstacle(world_table.objects[0])
    world.add_obstacle(world_table.objects[1])
    if save_debug_data:
        world.save_world_as_mesh("scene.stl", process_color=False)
    robot_mesh = (
        WorldConfig.create_merged_mesh_world(world, process_color=False).mesh[0].get_trimesh_mesh()
    )

    mesh_dataset = MeshDataset(
        None,
        n_frames=n_frames,
        image_size=640,
        save_data_dir=None,
        trimesh_mesh=robot_mesh,
        fov_deg=fov_deg,
    )
    q_js = JointState(position=q, joint_names=kin_model.joint_names)

    return mesh_dataset, q_js


def mask_image(robot_file="ur5e.yml"):
    save_debug_data = False
    write_pointcloud = False
    # create robot segmenter:
    tensor_args = TensorDeviceType()

    curobo_segmenter = RobotSegmenter.from_robot_file(
        robot_file, collision_sphere_buffer=0.01, distance_threshold=0.05, use_cuda_graph=True
    )

    mesh_dataset, q_js = create_render_dataset(robot_file, write_pointcloud, n_frames=20)

    if save_debug_data:
        visualize_scale = 10.0
        data = mesh_dataset[0]
        cam_obs = CameraObservation(
            depth_image=tensor_args.to_device(data["depth"]).unsqueeze(0) * 1000,
            intrinsics=data["intrinsics"],
            pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
        )
        # save depth image
        imageio.imwrite(
            "camera_depth.png",
            (cam_obs.depth_image * visualize_scale)
            .squeeze()
            .detach()
            .cpu()
            .numpy()
            .astype(np.uint16),
        )

        # save robot spheres in current joint configuration
        robot_kinematics = curobo_segmenter._robot_world.kinematics
        if write_pointcloud:
            sph = robot_kinematics.get_robot_as_spheres(q_js.position)
            WorldConfig(sphere=sph[0]).save_world_as_mesh("robot_spheres.stl")

            # save world pointcloud in robot origin

            pc = cam_obs.get_pointcloud()
            pc_obs = PointCloud("world", pose=cam_obs.pose.to_list(), points=pc)
            pc_obs.save_as_mesh("camera_pointcloud.stl", transform_with_pose=True)

        # run segmentation:
        depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
            cam_obs,
            q_js,
        )
        # save robot points as mesh

        robot_mask = cam_obs.clone()
        robot_mask.depth_image[~depth_mask] = 0.0

        if write_pointcloud:
            robot_mesh = PointCloud(
                "world", pose=robot_mask.pose.to_list(), points=robot_mask.get_pointcloud()
            )
            robot_mesh.save_as_mesh("robot_segmented.stl", transform_with_pose=True)
        # save depth image
        imageio.imwrite(
            "robot_depth.png",
            (robot_mask.depth_image * visualize_scale)
            .detach()
            .squeeze()
            .cpu()
            .numpy()
            .astype(np.uint16),
        )

        # save world points as mesh

        world_mask = cam_obs.clone()
        world_mask.depth_image[depth_mask] = 0.0
        if write_pointcloud:
            world_mesh = PointCloud(
                "world", pose=world_mask.pose.to_list(), points=world_mask.get_pointcloud()
            )
            world_mesh.save_as_mesh("world_segmented.stl", transform_with_pose=True)

        imageio.imwrite(
            "world_depth.png",
            (world_mask.depth_image * visualize_scale)
            .detach()
            .squeeze()
            .cpu()
            .numpy()
            .astype(np.uint16),
        )

    dt_list = []

    for i in range(len(mesh_dataset)):

        data = mesh_dataset[i]
        cam_obs = CameraObservation(
            depth_image=tensor_args.to_device(data["depth"]).unsqueeze(0) * 1000,
            intrinsics=data["intrinsics"],
            pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
        )
        if not curobo_segmenter.ready:
            curobo_segmenter.update_camera_projection(cam_obs)
        st_time = time.time()

        depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
            cam_obs,
            q_js,
        )

        torch.cuda.synchronize()
        dt_list.append(time.time() - st_time)

    print("Segmentation Time (ms), (hz)", np.mean(dt_list[5:]) * 1000.0, 1.0 / np.mean(dt_list[5:]))


def batch_mask_image(robot_file="ur5e.yml"):
    """Mask images from different camera views using batched query.

    Note: This only works for a single joint configuration across camera views.

    Args:
        robot_file: robot to use for example.
    """
    save_debug_data = True
    # create robot segmenter:
    tensor_args = TensorDeviceType()

    curobo_segmenter = RobotSegmenter.from_robot_file(
        robot_file, collision_sphere_buffer=0.01, distance_threshold=0.05, use_cuda_graph=True
    )

    mesh_dataset, q_js = create_render_dataset(robot_file, save_debug_data, fov_deg=60)

    mesh_dataset_zoom, q_js = create_render_dataset(
        robot_file, save_debug_data, fov_deg=40, n_frames=30
    )

    if save_debug_data:
        visualize_scale = 10.0
        data = mesh_dataset[0]
        cam_obs = CameraObservation(
            depth_image=tensor_args.to_device(data["depth"]).unsqueeze(0) * 1000,
            intrinsics=data["intrinsics"],
            pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
        )
        data_zoom = mesh_dataset_zoom[1]
        cam_obs = CameraObservation(
            depth_image=tensor_args.to_device(data["depth"]) * 1000,
            intrinsics=data["intrinsics"],
            pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
        )
        cam_obs_zoom = CameraObservation(
            depth_image=tensor_args.to_device(data_zoom["depth"]) * 1000,
            intrinsics=data_zoom["intrinsics"],
            pose=Pose.from_matrix(data_zoom["pose"].to(device=tensor_args.device)),
        )
        cam_obs = cam_obs.stack(cam_obs_zoom)

        for i in range(cam_obs.depth_image.shape[0]):
            # save depth image
            imageio.imwrite(
                "camera_depth_" + str(i) + ".png",
                (cam_obs.depth_image[i] * visualize_scale)
                .squeeze()
                .detach()
                .cpu()
                .numpy()
                .astype(np.uint16),
            )

        # run segmentation:
        depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
            cam_obs,
            q_js,
        )
        # save robot points as mesh

        robot_mask = cam_obs.clone()
        robot_mask.depth_image[~depth_mask] = 0.0
        for i in range(cam_obs.depth_image.shape[0]):
            # save depth image
            imageio.imwrite(
                "robot_depth_" + str(i) + ".png",
                (robot_mask.depth_image[i] * visualize_scale)
                .detach()
                .squeeze()
                .cpu()
                .numpy()
                .astype(np.uint16),
            )

            # save world points as mesh

            imageio.imwrite(
                "world_depth_" + str(i) + ".png",
                (filtered_image[i] * visualize_scale)
                .detach()
                .squeeze()
                .cpu()
                .numpy()
                .astype(np.uint16),
            )

    dt_list = []

    for i in range(len(mesh_dataset)):

        data = mesh_dataset[i]
        data_zoom = mesh_dataset_zoom[i + 1]
        cam_obs = CameraObservation(
            depth_image=tensor_args.to_device(data["depth"]) * 1000,
            intrinsics=data["intrinsics"],
            pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
        )
        cam_obs_zoom = CameraObservation(
            depth_image=tensor_args.to_device(data_zoom["depth"]) * 1000,
            intrinsics=data_zoom["intrinsics"],
            pose=Pose.from_matrix(data_zoom["pose"].to(device=tensor_args.device)),
        )
        cam_obs = cam_obs.stack(cam_obs_zoom)
        if not curobo_segmenter.ready:
            curobo_segmenter.update_camera_projection(cam_obs)
        st_time = time.time()

        depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
            cam_obs,
            q_js,
        )

        torch.cuda.synchronize()
        dt_list.append(time.time() - st_time)

    print("Segmentation Time (ms), (hz)", np.mean(dt_list[5:]) * 1000.0, 1.0 / np.mean(dt_list[5:]))


def batch_robot_mask_image(robot_file="ur5e.yml"):
    """Mask images from different camera views using batched query.

    Note: This example treats each image to have different robot joint configuration

    Args:
        robot_file: robot to use for example.
    """
    save_debug_data = True
    # create robot segmenter:
    tensor_args = TensorDeviceType()

    curobo_segmenter = RobotSegmenter.from_robot_file(
        robot_file, collision_sphere_buffer=0.01, distance_threshold=0.05, use_cuda_graph=True
    )

    mesh_dataset, q_js = create_render_dataset(robot_file, save_debug_data, fov_deg=60)

    mesh_dataset_zoom, q_js_zoom = create_render_dataset(
        robot_file, save_debug_data, fov_deg=60, retract_delta=-0.5
    )
    q_js = q_js.unsqueeze(0)
    q_js = q_js.stack(q_js_zoom.unsqueeze(0))

    if save_debug_data:
        visualize_scale = 10.0
        data = mesh_dataset[0]
        cam_obs = CameraObservation(
            depth_image=tensor_args.to_device(data["depth"]).unsqueeze(0) * 1000,
            intrinsics=data["intrinsics"],
            pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
        )
        data_zoom = mesh_dataset_zoom[0]
        cam_obs = CameraObservation(
            depth_image=tensor_args.to_device(data["depth"]) * 1000,
            intrinsics=data["intrinsics"],
            pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
        )
        cam_obs_zoom = CameraObservation(
            depth_image=tensor_args.to_device(data_zoom["depth"]) * 1000,
            intrinsics=data_zoom["intrinsics"],
            pose=Pose.from_matrix(data_zoom["pose"].to(device=tensor_args.device)),
        )
        cam_obs = cam_obs.stack(cam_obs_zoom)

        for i in range(cam_obs.depth_image.shape[0]):
            # save depth image
            imageio.imwrite(
                "camera_depth_" + str(i) + ".png",
                (cam_obs.depth_image[i] * visualize_scale)
                .squeeze()
                .detach()
                .cpu()
                .numpy()
                .astype(np.uint16),
            )

        # run segmentation:
        depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
            cam_obs,
            q_js,
        )
        # save robot points as mesh

        robot_mask = cam_obs.clone()
        robot_mask.depth_image[~depth_mask] = 0.0
        for i in range(cam_obs.depth_image.shape[0]):
            # save depth image
            imageio.imwrite(
                "robot_depth_" + str(i) + ".png",
                (robot_mask.depth_image[i] * visualize_scale)
                .detach()
                .squeeze()
                .cpu()
                .numpy()
                .astype(np.uint16),
            )

            # save world points as mesh

            imageio.imwrite(
                "world_depth_" + str(i) + ".png",
                (filtered_image[i] * visualize_scale)
                .detach()
                .squeeze()
                .cpu()
                .numpy()
                .astype(np.uint16),
            )

    dt_list = []

    for i in range(len(mesh_dataset)):

        data = mesh_dataset[i]
        data_zoom = mesh_dataset_zoom[i]
        cam_obs = CameraObservation(
            depth_image=tensor_args.to_device(data["depth"]) * 1000,
            intrinsics=data["intrinsics"],
            pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
        )
        cam_obs_zoom = CameraObservation(
            depth_image=tensor_args.to_device(data_zoom["depth"]) * 1000,
            intrinsics=data_zoom["intrinsics"],
            pose=Pose.from_matrix(data_zoom["pose"].to(device=tensor_args.device)),
        )
        cam_obs = cam_obs.stack(cam_obs_zoom)
        if not curobo_segmenter.ready:
            curobo_segmenter.update_camera_projection(cam_obs)
        st_time = time.time()

        depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
            cam_obs,
            q_js,
        )

        torch.cuda.synchronize()
        dt_list.append(time.time() - st_time)

    print("Segmentation Time (ms), (hz)", np.mean(dt_list[5:]) * 1000.0, 1.0 / np.mean(dt_list[5:]))


def profile_mask_image(robot_file="ur5e.yml"):
    # create robot segmenter:
    tensor_args = TensorDeviceType()

    curobo_segmenter = RobotSegmenter.from_robot_file(
        robot_file, collision_sphere_buffer=0.0, distance_threshold=0.05, use_cuda_graph=False
    )

    mesh_dataset, q_js = create_render_dataset(robot_file)

    dt_list = []
    data = mesh_dataset[0]
    cam_obs = CameraObservation(
        depth_image=tensor_args.to_device(data["depth"]).unsqueeze(0) * 1000,
        intrinsics=data["intrinsics"],
        pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
    )
    if not curobo_segmenter.ready:
        curobo_segmenter.update_camera_projection(cam_obs)
    depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(cam_obs, q_js)

    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:

        for i in range(len(mesh_dataset)):
            with profiler.record_function("get_data"):

                data = mesh_dataset[i]
                cam_obs = CameraObservation(
                    depth_image=tensor_args.to_device(data["depth"]).unsqueeze(0) * 1000,
                    intrinsics=data["intrinsics"],
                    pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
                )
            st_time = time.time()
            with profiler.record_function("segmentation"):

                depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
                    cam_obs, q_js
                )

    print("Exporting the trace..")
    prof.export_chrome_trace("segmentation.json")


if __name__ == "__main__":
    mask_image()
    # profile_mask_image()
    # batch_mask_image()
    # batch_robot_mask_image()