# Copyright (c) 2022-2023, The ORBIT Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

"""
Script to run an environment with a pick and lift state machine.

The state machine is implemented in the kernel function `infer_state_machine`.
It uses the `warp` library to run the state machine in parallel on the GPU.

Usage:
    python play_lift.py --num_envs 128
"""

from __future__ import annotations

"""Launch Omniverse Toolkit first."""

import argparse

from omni.isaac.orbit.app import AppLauncher

# add argparse arguments
parser = argparse.ArgumentParser(description="Pick and lift state machine for lift environments.")
parser.add_argument("--cpu", action="store_true", default=False, help="Use CPU pipeline.")
parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
# append AppLauncher cli args
AppLauncher.add_app_launcher_args(parser)
# parse the arguments
args_cli = parser.parse_args()

# launch omniverse app
app_launcher = AppLauncher(headless=args_cli.headless)
simulation_app = app_launcher.app

"""Rest everything else."""

import gymnasium as gym
import torch
import traceback
from enum import Enum
from typing import Sequence

import carb
import warp as wp

from omni.isaac.orbit.utils.timer import Timer

import omni.isaac.orbit_tasks  # noqa: F401
from omni.isaac.orbit_tasks.utils.parse_cfg import parse_env_cfg

# initialize warp
wp.init()


class GripperState(Enum):
    """States for the gripper."""

    OPEN = wp.constant(1.0)
    CLOSE = wp.constant(-1.0)


class PickSmState(Enum):
    """States for the pick state machine."""

    REST = wp.constant(0)
    APPROACH_ABOVE_OBJECT = wp.constant(1)
    APPROACH_OBJECT = wp.constant(2)
    GRASP_OBJECT = wp.constant(3)
    LIFT_OBJECT = wp.constant(4)
    DROP_OBJECT = wp.constant(5)


class PickSmWaitTime(Enum):
    """Additional wait times (in s) for states for before switching."""

    REST = wp.constant(1.0)
    APPROACH_ABOVE_OBJECT = wp.constant(0.5)
    APPROACH_OBJECT = wp.constant(0.3)
    GRASP_OBJECT = wp.constant(0.3)
    LIFT_OBJECT = wp.constant(2.0)
    DROP_OBJECT = wp.constant(0.2)


@wp.kernel
def infer_state_machine(
    dt: wp.array(dtype=wp.float32),
    sm_state: wp.array(dtype=wp.int32),
    sm_wait_time: wp.array(dtype=wp.float32),
    ee_pose: wp.array(dtype=wp.transform),
    object_pose: wp.array(dtype=wp.transform),
    des_object_pose: wp.array(dtype=wp.transform),
    des_ee_pose: wp.array(dtype=wp.transform),
    gripper_state: wp.array(dtype=wp.float32),
    offset: wp.array(dtype=wp.transform),
):
    # retrieve thread id
    tid = wp.tid()
    # retrieve state machine state
    state = sm_state[tid]
    # decide next state
    if state == PickSmState.REST.value:
        des_ee_pose[tid] = ee_pose[tid]
        gripper_state[tid] = GripperState.OPEN.value
        # wait for a while
        if sm_wait_time[tid] >= PickSmWaitTime.REST.value:
            # move to next state and reset wait time
            sm_state[tid] = PickSmState.APPROACH_ABOVE_OBJECT.value
            sm_wait_time[tid] = 0.0
    elif state == PickSmState.APPROACH_ABOVE_OBJECT.value:
        des_ee_pose[tid] = object_pose[tid]
        # TODO: This is causing issues.
        # des_ee_pose[tid] = wp.transform_multiply(des_ee_pose[tid], offset[tid])
        gripper_state[tid] = GripperState.OPEN.value
        # TODO: error between current and desired ee pose below threshold
        # wait for a while
        if sm_wait_time[tid] >= PickSmWaitTime.APPROACH_OBJECT.value:
            # move to next state and reset wait time
            sm_state[tid] = PickSmState.APPROACH_OBJECT.value
            sm_wait_time[tid] = 0.0
    elif state == PickSmState.APPROACH_OBJECT.value:
        des_ee_pose[tid] = object_pose[tid]
        gripper_state[tid] = GripperState.OPEN.value
        # TODO: error between current and desired ee pose below threshold
        # wait for a while
        if sm_wait_time[tid] >= PickSmWaitTime.APPROACH_OBJECT.value:
            # move to next state and reset wait time
            sm_state[tid] = PickSmState.GRASP_OBJECT.value
            sm_wait_time[tid] = 0.0
    elif state == PickSmState.GRASP_OBJECT.value:
        des_ee_pose[tid] = object_pose[tid]
        gripper_state[tid] = GripperState.CLOSE.value
        # wait for a while
        if sm_wait_time[tid] >= PickSmWaitTime.GRASP_OBJECT.value:
            # move to next state and reset wait time
            sm_state[tid] = PickSmState.LIFT_OBJECT.value
            sm_wait_time[tid] = 0.0
    elif state == PickSmState.LIFT_OBJECT.value:
        des_ee_pose[tid] = des_object_pose[tid]
        gripper_state[tid] = GripperState.CLOSE.value
        # TODO: error between current and desired ee pose below threshold
        # wait for a while
        if sm_wait_time[tid] >= PickSmWaitTime.LIFT_OBJECT.value:
            # move to next state and reset wait time
            sm_state[tid] = PickSmState.DROP_OBJECT.value
            sm_wait_time[tid] = 0.0
    elif state == PickSmState.DROP_OBJECT.value:
        des_ee_pose[tid] = des_object_pose[tid]
        gripper_state[tid] = GripperState.OPEN.value
        # wait for a while
        if sm_wait_time[tid] >= PickSmWaitTime.DROP_OBJECT.value:
            # move to next state and reset wait time
            sm_state[tid] = PickSmState.REST.value
            sm_wait_time[tid] = 0.0
    # increment wait time
    sm_wait_time[tid] = sm_wait_time[tid] + dt[tid]


class PickAndLiftSm:
    """A simple state machine in a robot's task space to pick and lift an object.

    The state machine is implemented as a warp kernel. It takes in the current state of
    the robot's end-effector and the object, and outputs the desired state of the robot's
    end-effector and the gripper. The state machine is implemented as a finite state
    machine with the following states:

    1. REST: The robot is at rest.
    2. APPROACH_ABOVE_OBJECT: The robot moves above the object.
    3. APPROACH_OBJECT: The robot moves to the object.
    4. GRASP_OBJECT: The robot grasps the object.
    5. LIFT_OBJECT: The robot lifts the object.
    6. DROP_OBJECT: The robot drops the object.
    """

    def __init__(self, dt: float, num_envs: int, device: torch.device | str = "cpu"):
        """Initialize the state machine.

        Args:
            dt: The environment time step.
            num_envs: The number of environments to simulate.
            device: The device to run the state machine on.
        """
        # save parameters
        self.dt = dt
        self.num_envs = num_envs
        self.device = device
        # initialize state machine
        self.sm_dt = torch.full((self.num_envs,), self.dt, dtype=torch.float32, device=self.device)
        self.sm_state = torch.full((self.num_envs,), 0, dtype=torch.int32, device=self.device)
        self.sm_wait_time = torch.zeros((self.num_envs,), dtype=torch.float32, device=self.device)
        # desired state
        self.des_ee_pose = torch.zeros((self.num_envs, 7), dtype=torch.float32, device=self.device)
        self.des_gripper_state = torch.full((self.num_envs,), 0, dtype=torch.float32, device=self.device)
        # approach above object offset
        self.offset = torch.zeros((self.num_envs, 7), dtype=torch.float32, device=self.device)
        self.offset[:, 2] = 0.1
        # convert to warp
        self.sm_dt_wp = wp.from_torch(self.sm_dt, wp.float32)
        self.sm_state_wp = wp.from_torch(self.sm_state, wp.int32)
        self.sm_wait_time_wp = wp.from_torch(self.sm_wait_time, wp.float32)
        self.des_ee_pose_wp = wp.from_torch(self.des_ee_pose, wp.transform)
        self.des_gripper_state_wp = wp.from_torch(self.des_gripper_state, wp.float32)
        self.offset_wp = wp.from_torch(self.offset, wp.transform)

    def reset_idx(self, env_ids: Sequence[int] = None):
        """Reset the state machine."""
        if env_ids is None:
            env_ids = slice(None)
        self.sm_state[env_ids] = 0
        self.sm_wait_time[env_ids] = 0.0

    def compute(self, ee_pose: torch.Tensor, object_pose: torch.Tensor, des_object_pose: torch.Tensor):
        """Compute the desired state of the robot's end-effector and the gripper."""
        # convert to warp
        ee_pose_wp = wp.from_torch(ee_pose.contiguous(), wp.transform)
        object_pose_wp = wp.from_torch(object_pose.contiguous(), wp.transform)
        des_object_pose_wp = wp.from_torch(des_object_pose.contiguous(), wp.transform)
        # run state machine
        wp.launch(
            kernel=infer_state_machine,
            dim=self.num_envs,
            inputs=[
                self.sm_dt_wp,
                self.sm_state_wp,
                self.sm_wait_time_wp,
                ee_pose_wp,
                object_pose_wp,
                des_object_pose_wp,
                self.des_ee_pose_wp,
                self.des_gripper_state_wp,
                self.offset_wp,
            ],
        )
        wp.synchronize()
        # convert to torch
        return torch.cat([self.des_ee_pose, self.des_gripper_state.unsqueeze(-1)], dim=-1)


def main():
    # parse configuration
    env_cfg = parse_env_cfg("Isaac-Lift-Franka-v0", use_gpu=not args_cli.cpu, num_envs=args_cli.num_envs)
    # -- control configuration
    env_cfg.control.control_type = "inverse_kinematics"
    env_cfg.control.inverse_kinematics.command_type = "position_abs"
    # -- randomization configuration
    env_cfg.randomization.object_initial_pose.position_cat = "uniform"
    env_cfg.randomization.object_desired_pose.position_cat = "uniform"
    # -- termination configuration
    env_cfg.terminations.episode_timeout = False
    # -- robot configuration
    env_cfg.robot.robot_type = "franka"
    # create environment
    env = gym.make("Isaac-Lift-Franka-v0", cfg=env_cfg)

    # create action buffers
    actions = torch.zeros((env.num_envs, env.action_space.shape[0]), device=env.device)
    # create state machine
    pick_sm = PickAndLiftSm(env.dt, env.num_envs, env.device)

    # reset environment
    env.reset()

    # run state machine
    for _ in range(10000):
        # run everything in inference mode
        with torch.inference_mode():
            # step environment
            dones = env.step(actions)[-2]

            # observations
            ee_pose = env.robot.data.ee_state_w[:, :7].clone()
            object_pose = env.object.data.root_state_w[:, :7].clone()
            des_object_pose = env.object_des_pose_w.clone()
            # transform from world to base frames
            ee_pose[:, :3] -= env.robot.data.root_pos_w
            object_pose[:, :3] -= env.robot.data.root_pos_w
            des_object_pose[:, :3] -= env.robot.data.root_pos_w
            # advance state machine
            with Timer("state machine"):
                sm_actions = pick_sm.compute(ee_pose, object_pose, des_object_pose)

            # set actions for IK with positions
            actions[:, :3] = sm_actions[:, :3]
            actions[:, -1] = sm_actions[:, -1]
            # reset state machine
            if dones.any():
                pick_sm.reset_idx(dones.nonzero(as_tuple=False).squeeze(-1))


if __name__ == "__main__":
    try:
        # run the main execution
        main()
    except Exception as err:
        carb.log_error(err)
        carb.log_error(traceback.format_exc())
        raise
    finally:
        # close sim app
        simulation_app.close()
