#
# 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.
#

# Standard Library

# Third Party
import torch

# CuRobo
from curobo.geom.sdf.world import CollisionCheckerType
from curobo.geom.types import Cuboid, WorldConfig
from curobo.types.base import TensorDeviceType
from curobo.types.math import Pose
from curobo.types.robot import JointState, RobotConfig
from curobo.util.logger import setup_curobo_logger
from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig


def plot_traj(trajectory, dt):
    # Third Party
    import matplotlib.pyplot as plt

    _, axs = plt.subplots(4, 1)
    q = trajectory.position.cpu().numpy()
    qd = trajectory.velocity.cpu().numpy()
    qdd = trajectory.acceleration.cpu().numpy()
    qddd = trajectory.jerk.cpu().numpy()
    timesteps = [i * dt for i in range(q.shape[0])]
    for i in range(q.shape[-1]):
        axs[0].plot(timesteps, q[:, i], label=str(i))
        axs[1].plot(timesteps, qd[:, i], label=str(i))
        axs[2].plot(timesteps, qdd[:, i], label=str(i))
        axs[3].plot(timesteps, qddd[:, i], label=str(i))

    plt.legend()
    # plt.savefig("test.png")
    plt.show()


def plot_iters_traj(trajectory, d_id=1, dof=7, seed=0):
    # Third Party
    import matplotlib.pyplot as plt

    _, axs = plt.subplots(len(trajectory), 1)
    if len(trajectory) == 1:
        axs = [axs]
    for k in range(len(trajectory)):
        q = trajectory[k]

        for i in range(len(q)):
            axs[k].plot(
                q[i][seed, :-1, d_id].cpu(),
                "r+-",
                label=str(i),
                alpha=0.1 + min(0.9, float(i) / (len(q))),
            )
    plt.legend()
    plt.show()


def plot_iters_traj_3d(trajectory, d_id=1, dof=7, seed=0):
    # Third Party
    import matplotlib.pyplot as plt

    ax = plt.axes(projection="3d")
    c = 0
    h = trajectory[0][0].shape[1] - 1
    x = [x for x in range(h)]

    for k in range(len(trajectory)):
        q = trajectory[k]

        for i in range(len(q)):
            # ax.plot3D(x,[c for _ in range(h)],  q[i][seed, :, d_id].cpu())#, 'r')
            ax.scatter3D(
                x, [c for _ in range(h)], q[i][seed, :h, d_id].cpu(), c=q[i][seed, :, d_id].cpu()
            )
            # @plt.show()
            c += 1
    # plt.legend()
    plt.show()


def demo_motion_gen_mesh():
    PLOT = False
    tensor_args = TensorDeviceType()
    world_file = "collision_mesh_scene.yml"
    robot_file = "franka.yml"
    motion_gen_config = MotionGenConfig.load_from_robot_config(
        robot_file,
        world_file,
        tensor_args,
        # trajopt_tsteps=40,
        collision_checker_type=CollisionCheckerType.MESH,
        use_cuda_graph=False,
    )
    motion_gen = MotionGen(motion_gen_config)
    robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
    robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
    retract_cfg = robot_cfg.retract_config
    state = motion_gen.rollout_fn.compute_kinematics(
        JointState.from_position(retract_cfg.view(1, -1))
    )

    retract_pose = Pose(state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze())
    start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.5)
    result = motion_gen.plan(
        start_state,
        retract_pose,
        enable_graph=False,
        enable_opt=True,
        max_attempts=1,
        num_trajopt_seeds=10,
        num_graph_seeds=10,
    )
    print(result.status, result.attempts, result.trajopt_time)
    traj = result.raw_plan  # optimized plan
    print("Trajectory Generated: ", result.success)
    if PLOT:
        plot_traj(traj.cpu().numpy())


def demo_motion_gen(js=False):
    # Standard Library
    PLOT = True
    tensor_args = TensorDeviceType()
    world_file = "collision_table.yml"
    robot_file = "franka.yml"
    motion_gen_config = MotionGenConfig.load_from_robot_config(
        robot_file,
        world_file,
        tensor_args,
        interpolation_dt=0.01,
    )

    motion_gen = MotionGen(motion_gen_config)

    motion_gen.warmup(enable_graph=True, warmup_js_trajopt=js)
    robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
    robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
    retract_cfg = motion_gen.get_retract_config()
    state = motion_gen.rollout_fn.compute_kinematics(
        JointState.from_position(retract_cfg.view(1, -1))
    )

    retract_pose = Pose(state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze())
    start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.1)
    goal_state = start_state.clone()
    goal_state.position[..., 3] -= 0.1
    if js:
        result = motion_gen.plan_single_js(
            start_state,
            goal_state,
            MotionGenPlanConfig(
                max_attempts=1, enable_graph=False, enable_opt=True, enable_finetune_trajopt=True
            ),
        )
    else:
        result = motion_gen.plan_single(
            start_state, retract_pose, MotionGenPlanConfig(max_attempts=1)
        )
    traj = result.get_interpolated_plan()
    print("Trajectory Generated: ", result.success, result.solve_time, result.status)
    if PLOT and result.success.item():
        plot_traj(traj, result.interpolation_dt)
        # plt.save("test.png")
        # plt.close()
        # traj = result.debug_info["opt_solution"].position.view(-1,7)
        # plot_traj(traj.cpu().numpy())


def demo_motion_gen_debug():
    PLOT = True
    tensor_args = TensorDeviceType()
    world_file = "collision_cubby.yml"
    robot_file = "franka.yml"
    motion_gen_config = MotionGenConfig.load_from_robot_config(
        robot_file,
        world_file,
        tensor_args,
        trajopt_tsteps=24,
        collision_checker_type=CollisionCheckerType.PRIMITIVE,
        use_cuda_graph=True,
        num_trajopt_seeds=1,
        num_graph_seeds=1,
        store_ik_debug=True,
        store_trajopt_debug=True,
        trajopt_particle_opt=False,
        grad_trajopt_iters=100,
    )
    motion_gen = MotionGen(motion_gen_config)
    robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
    robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
    retract_cfg = robot_cfg.cspace.retract_config
    state = motion_gen.rollout_fn.compute_kinematics(
        JointState.from_position(retract_cfg.view(1, -1))
    )

    retract_pose = Pose(state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze())
    start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.4)
    result = motion_gen.plan(start_state, retract_pose, enable_graph=True, enable_opt=True)
    if result.status not in [None, "Opt Fail"]:
        return
    traj = result.plan.view(-1, 7)  # optimized plan
    print("Trajectory Generated: ", result.success)
    trajectory_iter_steps = result.debug_info["trajopt_result"].debug_info["solver"]["steps"]

    if PLOT:
        plot_iters_traj_3d(trajectory_iter_steps, d_id=6)
        # plot_traj(traj.cpu().numpy())


def demo_motion_gen_batch():
    PLOT = False
    tensor_args = TensorDeviceType()
    world_file = "collision_cubby.yml"
    robot_file = "franka.yml"
    motion_gen_config = MotionGenConfig.load_from_robot_config(
        robot_file,
        world_file,
        tensor_args,
        collision_checker_type=CollisionCheckerType.PRIMITIVE,
        use_cuda_graph=True,
        num_trajopt_seeds=12,
        num_graph_seeds=1,
        num_ik_seeds=30,
    )
    motion_gen = MotionGen(motion_gen_config)
    robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
    robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
    retract_cfg = motion_gen.get_retract_config()
    state = motion_gen.rollout_fn.compute_kinematics(
        JointState.from_position(retract_cfg.view(1, -1))
    )

    retract_pose = Pose(state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze())
    start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.6)

    retract_pose = retract_pose.repeat_seeds(2)
    retract_pose.position[0, 0] = -0.3
    result = motion_gen.plan_batch(
        start_state.repeat_seeds(2),
        retract_pose,
        MotionGenPlanConfig(
            max_attempts=5, enable_graph=False, enable_graph_attempt=1, enable_opt=True
        ),
    )
    traj = result.optimized_plan.position.view(2, -1, 7)  # optimized plan
    print("Trajectory Generated: ", result.success)
    if PLOT:
        plot_traj(traj[0, : result.path_buffer_last_tstep[0], :].cpu().numpy())
        plot_traj(traj[1, : result.path_buffer_last_tstep[1], :].cpu().numpy())


def demo_motion_gen_api():
    tensor_args = TensorDeviceType(device=torch.device("cuda:0"))
    interpolation_dt = 0.02
    # create motion gen with a cuboid cache to be able to load obstacles later:
    motion_gen_cfg = MotionGenConfig.load_from_robot_config(
        "franka.yml",
        "collision_table.yml",
        tensor_args,
        trajopt_tsteps=34,
        interpolation_steps=5000,
        num_ik_seeds=50,
        num_trajopt_seeds=6,
        collision_checker_type=CollisionCheckerType.PRIMITIVE,
        grad_trajopt_iters=500,
        trajopt_dt=0.5,
        interpolation_dt=interpolation_dt,
        evaluate_interpolated_trajectory=True,
        js_trajopt_dt=0.5,
        js_trajopt_tsteps=34,
    )
    motion_gen = MotionGen(motion_gen_cfg)

    motion_gen.warmup(warmup_js_trajopt=False)

    # create world representation:

    cuboids = [Cuboid(name="obs_1", pose=[0, 0, 0, 1, 0, 0, 0], dims=[0.1, 0.1, 0.1])]
    world = WorldConfig(cuboid=cuboids)

    motion_gen.update_world(world)

    q_start = JointState.from_position(
        tensor_args.to_device([[0.0, -1.3, 0.0, -2.5, 0.0, 1.0, 0.0]]),
        joint_names=[
            "panda_joint1",
            "panda_joint2",
            "panda_joint3",
            "panda_joint4",
            "panda_joint5",
            "panda_joint6",
            "panda_joint7",
        ],
    )

    goal_pose = Pose(
        position=tensor_args.to_device([[0.5, 0.0, 0.3]]),
        quaternion=tensor_args.to_device([[1, 0, 0, 0]]),
    )
    print(goal_pose)

    result = motion_gen.plan_single(q_start, goal_pose)

    # get result:

    interpolated_solution = result.get_interpolated_plan()

    print(result.get_interpolated_plan())


def demo_motion_gen_batch_env(n_envs: int = 10):
    tensor_args = TensorDeviceType()
    world_files = ["collision_table.yml" for _ in range(int(n_envs / 2))] + [
        "collision_test.yml" for _ in range(int(n_envs / 2))
    ]

    world_cfg = [
        WorldConfig.from_dict(load_yaml(join_path(get_world_configs_path(), world_file)))
        for world_file in world_files
    ]
    robot_file = "franka.yml"
    motion_gen_config = MotionGenConfig.load_from_robot_config(
        robot_file,
        world_cfg,
        tensor_args,
        trajopt_tsteps=30,
        use_cuda_graph=False,
        num_trajopt_seeds=4,
        num_ik_seeds=30,
        num_batch_ik_seeds=30,
        evaluate_interpolated_trajectory=True,
        interpolation_dt=0.05,
        interpolation_steps=500,
        grad_trajopt_iters=30,
    )
    motion_gen_batch_env = MotionGen(motion_gen_config)
    motion_gen_batch_env.reset()
    motion_gen_batch_env.warmup(
        enable_graph=False, batch=n_envs, warmup_js_trajopt=False, batch_env_mode=True
    )
    retract_cfg = motion_gen_batch_env.get_retract_config().clone()
    state = motion_gen_batch_env.compute_kinematics(
        JointState.from_position(retract_cfg.view(1, -1))
    )

    goal_pose = Pose(
        state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze()
    ).repeat_seeds(n_envs)

    start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3).repeat_seeds(n_envs)

    goal_pose.position[1, 0] -= 0.2

    m_config = MotionGenPlanConfig(
        False, True, max_attempts=1, enable_graph_attempt=None, enable_finetune_trajopt=False
    )
    result = motion_gen_batch_env.plan_batch_env(start_state, goal_pose, m_config)

    print(n_envs, result.total_time, result.total_time / n_envs)


if __name__ == "__main__":
    setup_curobo_logger("error")
    # demo_motion_gen(js=True)
    demo_motion_gen_batch()
    # n = [2, 10]
    # for n_envs in n:
    #    demo_motion_gen_batch_env(n_envs=n_envs)
