# Copyright (c) 2024-2025, The Isaac Lab Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0

"""
Script to add mimic annotations to demos to be used as source demos for mimic dataset generation.
"""

# Launching Isaac Sim Simulator first.

import argparse

from omni.isaac.lab.app import AppLauncher

# add argparse arguments
parser = argparse.ArgumentParser(description="Annotate demonstrations for Isaac Lab environments.")
parser.add_argument("--task", type=str, default=None, help="Name of the task.")
parser.add_argument(
    "--input_file", type=str, default="./datasets/dataset.hdf5", help="File name of the dataset to be annotated."
)
parser.add_argument(
    "--output_file",
    type=str,
    default="./datasets/dataset_annotated.hdf5",
    help="File name of the annotated output dataset file.",
)
parser.add_argument("--auto", action="store_true", default=False, help="Automatically annotate subtasks.")
parser.add_argument(
    "--signals",
    type=str,
    nargs="+",
    default=[],
    help="Sequence of subtask termination signals for all except last subtask",
)

# append AppLauncher cli args
AppLauncher.add_app_launcher_args(parser)
# parse the arguments
args_cli = parser.parse_args()

# launch the simulator
app_launcher = AppLauncher(args_cli)
simulation_app = app_launcher.app

"""Rest everything follows."""

import contextlib
import gymnasium as gym
import os
import torch

import omni.isaac.lab_mimic.envs

# Only enables inputs if this script is NOT headless mode
if not args_cli.headless and not os.environ.get("HEADLESS", 0):
    from omni.isaac.lab.devices import Se3Keyboard
from omni.isaac.lab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg
from omni.isaac.lab.managers import RecorderTerm, RecorderTermCfg
from omni.isaac.lab.utils import configclass
from omni.isaac.lab.utils.datasets import HDF5DatasetFileHandler

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

is_paused = False
current_action_index = 0
subtask_indices = []


def play_cb():
    global is_paused
    is_paused = False


def pause_cb():
    global is_paused
    is_paused = True


def mark_subtask_cb():
    global current_action_index, subtask_indices
    subtask_indices.append(current_action_index)
    print(f"Marked subtask at action index: {current_action_index}")


class PreStepDatagenInfoRecorder(RecorderTerm):
    """Recorder term that records the datagen info data in each step."""

    def record_pre_step(self):
        datagen_info = {
            "object_pose": self._env.scene.get_state(is_relative=True)["rigid_object"],
            "target_eef_pose": self._env.action_to_target_eef_pos(self._env.action_manager.action),
        }
        return "obs", datagen_info


@configclass
class PreStepDatagenInfoRecorderCfg(RecorderTermCfg):
    """Configuration for the datagen info recorder term."""

    class_type: type[RecorderTerm] = PreStepDatagenInfoRecorder


class PreStepSubtaskTermsObservationsRecorder(RecorderTerm):
    """Recorder term that records the subtask completion observations in each step."""

    def record_pre_step(self):
        return "obs/subtask_term_signals", self._env.obs_buf["subtask_terms"]


@configclass
class PreStepSubtaskTermsObservationsRecorderCfg(RecorderTermCfg):
    """Configuration for the step subtask terms observation recorder term."""

    class_type: type[RecorderTerm] = PreStepSubtaskTermsObservationsRecorder


@configclass
class MimicRecorderManagerCfg(ActionStateRecorderManagerCfg):
    """Mimic specific recorder terms."""

    record_pre_step_datagen_info = PreStepDatagenInfoRecorderCfg()
    record_pre_step_subtask_term_signals = PreStepSubtaskTermsObservationsRecorderCfg()


def main():
    """Add Isaac Lab Mimic annotations to the given demo dataset file."""
    global is_paused, current_action_index, subtask_indices

    if not args_cli.auto and len(args_cli.signals) == 0:
        if len(args_cli.signals) == 0:
            raise ValueError("Subtask signals should be provided for manual mode.")

    # Load input dataset to be annotated
    if not os.path.exists(args_cli.input_file):
        raise FileNotFoundError(f"The input dataset file {args_cli.input_file} does not exist.")
    dataset_file_handler = HDF5DatasetFileHandler()
    dataset_file_handler.open(args_cli.input_file)
    env_name = dataset_file_handler.get_env_name()
    episode_count = dataset_file_handler.get_num_episodes()

    if episode_count == 0:
        print("No episodes found in the dataset.")
        exit()

    # get output directory path and file name (without extension) from cli arguments
    output_dir = os.path.dirname(args_cli.output_file)
    output_file_name = os.path.splitext(os.path.basename(args_cli.output_file))[0]
    # create output directory if it does not exist
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    if args_cli.task is not None:
        env_name = args_cli.task
    if env_name is None:
        raise ValueError("Task/env name was not specified nor found in the dataset.")

    env_cfg = parse_env_cfg(env_name, device=args_cli.device, num_envs=1)

    # Disable all termination terms
    env_cfg.terminations = {}

    # Set up recorder terms for mimic annotations
    env_cfg.env_name = args_cli.task
    env_cfg.recorders: MimicRecorderManagerCfg = MimicRecorderManagerCfg()
    env_cfg.recorders.dataset_export_dir_path = output_dir
    env_cfg.recorders.dataset_filename = output_file_name

    # create environment from loaded config
    env = gym.make(args_cli.task, cfg=env_cfg)

    # reset environment
    env.reset()

    # Only enables inputs if this script is NOT headless mode
    if not args_cli.headless and not os.environ.get("HEADLESS", 0):
        keyboard_interface = Se3Keyboard(pos_sensitivity=0.1, rot_sensitivity=0.1)
        keyboard_interface.add_callback("N", play_cb)
        keyboard_interface.add_callback("B", pause_cb)
        if not args_cli.auto:
            keyboard_interface.add_callback("S", mark_subtask_cb)
        keyboard_interface.reset()

    # simulate environment -- run everything in inference mode
    with contextlib.suppress(KeyboardInterrupt) and torch.inference_mode():
        while simulation_app.is_running() and not simulation_app.is_exiting():
            # Iterate over the episodes in the loaded dataset file
            for episode_index, episode_name in enumerate(dataset_file_handler.get_episode_names()):
                subtask_indices = []
                print(f"\nAnnotating episode #{episode_index} ({episode_name})")
                episode = dataset_file_handler.load_episode(episode_name, env.unwrapped.device)
                episode_data = episode.data

                # read initial state from the loaded episode
                initial_state = episode_data["initial_state"]
                env.unwrapped.recorder_manager.reset()
                env.unwrapped.reset_to(initial_state, None, is_relative=True)

                # replay actions from this episode
                actions = episode_data["actions"]
                first_action = True
                for action_index, action in enumerate(actions):
                    current_action_index = action_index
                    if first_action:
                        first_action = False
                    else:
                        while is_paused:
                            env.unwrapped.sim.render()
                            continue
                    action_tensor = torch.Tensor(action).reshape([1, action.shape[0]])
                    env.step(torch.Tensor(action_tensor))

                if not args_cli.auto:
                    print(f"\tSubtasks marked at action indices: {subtask_indices}")
                    if len(args_cli.signals) != len(subtask_indices):
                        raise ValueError(
                            f"Number of annotated subtask signals {len(subtask_indices)} should be equal               "
                            f"                          to number of subtasks {len(args_cli.signals)}"
                        )
                    annotated_episode = env.unwrapped.recorder_manager.get_episode(0)
                    del annotated_episode.data["obs"]["subtask_term_signals"]
                    for subtask_index in range(len(args_cli.signals)):
                        # subtask termination signal is false until subtask is complete, and true afterwards
                        subtask_signals = torch.ones(len(actions), dtype=torch.bool)
                        subtask_signals[: subtask_indices[subtask_index]] = False
                        annotated_episode.add(
                            f"obs/subtask_term_signals/{args_cli.signals[subtask_index]}", subtask_signals
                        )

                # set success to the recorded episode data and export to file
                env.unwrapped.recorder_manager.set_success_to_episodes(
                    None, torch.tensor([[True]], dtype=torch.bool, device=env.unwrapped.device)
                )
                env.unwrapped.recorder_manager.export_episodes()
                print("\tExported annotated episode.")

            break

    # Close environment after annotation is complete
    env.close()


if __name__ == "__main__":
    # run the main function
    main()
    # close sim app
    simulation_app.close()
