"""Script to play a checkpoint from an RSL-RL agent."""

"""Launch Isaac Sim Simulator first."""

import argparse
import sys
from pathlib import Path

from isaaclab.app import AppLauncher

# local imports
import cli_args  # isort: skip

# Prefer this checkout over any other editable ``engineai_lab`` installation.
REPO_ROOT = Path(__file__).resolve().parents[1]
SOURCE_ROOT = REPO_ROOT / "source"
if str(SOURCE_ROOT) not in sys.path:
    sys.path.insert(0, str(SOURCE_ROOT))

# add argparse arguments
parser = argparse.ArgumentParser(description="Play an RSL-RL policy checkpoint.")
parser.add_argument(
    "--disable_fabric", action="store_true", default=False, help="Disable fabric and use USD I/O operations."
)
parser.add_argument("--num_envs", type=int, default=None, help="Number of environments to simulate.")
parser.add_argument("--task", type=str, default=None, help="Name of the task.")
parser.add_argument("--max_steps", type=int, default=None, help="Stop replay after this many policy steps.")

command_group = parser.add_argument_group("velocity command", description="Velocity command source during replay.")
command_group.add_argument(
    "--command_source",
    type=str,
    choices=("random", "fixed", "keyboard"),
    default="random",
    help="Use environment-generated, fixed, or keyboard velocity commands.",
)
command_group.add_argument("--vx", type=float, default=0.3, help="Fixed forward velocity in m/s.")
command_group.add_argument("--vy", type=float, default=0.0, help="Fixed lateral velocity in m/s.")
command_group.add_argument("--wz", type=float, default=0.0, help="Fixed yaw velocity in rad/s.")
command_group.add_argument(
    "--command_mode",
    type=str,
    choices=("step", "ramp"),
    default="ramp",
    help="Apply external commands immediately or through acceleration limits.",
)
command_group.add_argument("--linear_accel", type=float, default=0.8, help="Planar acceleration limit in m/s^2.")
command_group.add_argument("--yaw_accel", type=float, default=1.5, help="Yaw acceleration limit in rad/s^2.")
command_group.add_argument("--keyboard_vx", type=float, default=0.4, help="Keyboard forward velocity in m/s.")
command_group.add_argument("--keyboard_vy", type=float, default=0.2, help="Keyboard lateral velocity in m/s.")
command_group.add_argument("--keyboard_wz", type=float, default=0.5, help="Keyboard yaw velocity in rad/s.")

# append RSL-RL cli arguments
cli_args.add_rsl_rl_args(parser)
# append AppLauncher cli args
AppLauncher.add_app_launcher_args(parser)
args_cli, hydra_args = parser.parse_known_args()
if args_cli.linear_accel <= 0.0:
    parser.error("--linear_accel must be positive.")
if args_cli.yaw_accel <= 0.0:
    parser.error("--yaw_accel must be positive.")
if args_cli.max_steps is not None and args_cli.max_steps <= 0:
    parser.error("--max_steps must be positive.")
# always enable cameras to record video
# clear out sys.argv for Hydra
sys.argv = [sys.argv[0]] + hydra_args

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

"""Rest everything follows."""

import os

import gymnasium as gym
import torch

from rsl_rl.runners import OnPolicyRunner

from isaaclab.envs import (
    DirectMARLEnv,
    DirectMARLEnvCfg,
    DirectRLEnvCfg,
    ManagerBasedRLEnvCfg,
    multi_agent_to_single_agent,
)
from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper
from isaaclab_tasks.utils import get_checkpoint_path
from isaaclab_tasks.utils.hydra import hydra_task_config

# Import extensions to set up environment tasks
import engineai_lab.tasks  # noqa: F401


class VelocityCommandController:
    """Apply fixed or keyboard commands directly to a velocity command term."""

    def __init__(self, env: RslRlVecEnvWrapper, args: argparse.Namespace):
        self.env = env
        self.args = args
        self.command_term = env.unwrapped.command_manager.get_term("base_velocity")
        if not hasattr(self.command_term, "vel_command_b"):
            raise TypeError("The base_velocity command term does not expose a velocity command buffer.")

        self.current_command = torch.zeros(3, device=env.unwrapped.device)
        self.fixed_command = torch.tensor([args.vx, args.vy, args.wz], device=env.unwrapped.device)
        self.keyboard = None

        if args.command_source == "keyboard":
            from isaaclab.devices import Se2Keyboard, Se2KeyboardCfg

            keyboard_cfg = Se2KeyboardCfg(
                sim_device=str(env.unwrapped.device),
                v_x_sensitivity=args.keyboard_vx,
                v_y_sensitivity=args.keyboard_vy,
                omega_z_sensitivity=args.keyboard_wz,
            )
            self.keyboard = Se2Keyboard(keyboard_cfg)
            print(self.keyboard)
        else:
            print(
                "[INFO] Fixed velocity command: "
                f"vx={args.vx:.3f} m/s, vy={args.vy:.3f} m/s, wz={args.wz:.3f} rad/s"
            )
        print(f"[INFO] External command mode: {args.command_mode}")

    def update(self):
        target = self.keyboard.advance() if self.keyboard is not None else self.fixed_command
        if self.args.command_mode == "step":
            self.current_command.copy_(target)
        else:
            self._apply_ramp(target)

        self.command_term.vel_command_b[:] = self.current_command.unsqueeze(0)
        self.command_term.time_left.fill_(float("inf"))
        self.command_term.is_standing_env.fill_(False)
        if hasattr(self.command_term, "is_heading_env"):
            self.command_term.is_heading_env.fill_(False)

    def _apply_ramp(self, target: torch.Tensor):
        dt = self.env.unwrapped.step_dt
        linear_delta = target[:2] - self.current_command[:2]
        linear_delta_norm = torch.linalg.norm(linear_delta)
        max_linear_delta = self.args.linear_accel * dt
        linear_scale = torch.clamp(max_linear_delta / torch.clamp(linear_delta_norm, min=1.0e-6), max=1.0)
        linear_delta = linear_delta * linear_scale
        self.current_command[:2] += linear_delta

        yaw_delta = torch.clamp(
            target[2] - self.current_command[2],
            min=-self.args.yaw_accel * dt,
            max=self.args.yaw_accel * dt,
        )
        self.current_command[2] += yaw_delta


@hydra_task_config(args_cli.task, "rsl_rl_cfg_entry_point")
def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlOnPolicyRunnerCfg):
    """Play with RSL-RL agent."""
    agent_cfg: RslRlOnPolicyRunnerCfg = cli_args.parse_rsl_rl_cfg(args_cli.task, args_cli)
    env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs

    # specify directory for logging experiments
    log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)
    log_root_path = os.path.abspath(log_root_path)
    print(f"[INFO] Loading experiment from directory: {log_root_path}")
    resume_path = get_checkpoint_path(log_root_path, agent_cfg.load_run, agent_cfg.load_checkpoint)

    log_dir = os.path.dirname(resume_path)

    # set the log directory for the environment (works for all environment types)
    env_cfg.log_dir = log_dir

    # Fixed/keyboard replay already applies its own step or ramp controller.
    # Disable the training-side ramp to avoid applying acceleration limits twice.
    if args_cli.command_source != "random" and hasattr(
        env_cfg.commands.base_velocity, "command_ramp_rates"
    ):
        env_cfg.commands.base_velocity.command_ramp_rates = None

    # create isaac environment
    env = gym.make(args_cli.task, cfg=env_cfg, render_mode=None)

    # convert to single-agent instance if required by the RL algorithm
    if isinstance(env.unwrapped, DirectMARLEnv):
        env = multi_agent_to_single_agent(env)

    # wrap around environment for rsl-rl
    env = RslRlVecEnvWrapper(env)

    command_controller = None
    if args_cli.command_source != "random":
        command_controller = VelocityCommandController(env, args_cli)

    # load previously trained model
    ppo_runner = OnPolicyRunner(env, agent_cfg.to_dict(), log_dir=None, device=agent_cfg.device)
    ppo_runner.load(resume_path)

    # obtain the trained policy for inference
    policy = ppo_runner.get_inference_policy(device=env.unwrapped.device)

    # export policy to onnx/jit
    export_model_dir = os.path.join(os.path.dirname(resume_path), "exported")
    file_basename = os.path.basename(resume_path).split(".")[0]
    ppo_runner.export_policy_to_onnx(path=export_model_dir, filename=file_basename+".onnx")

    # reset environment
    obs = env.get_observations()
    step_count = 0
    # simulate environment
    while simulation_app.is_running() and (args_cli.max_steps is None or step_count < args_cli.max_steps):
        # run everything in inference mode
        with torch.inference_mode():
            if command_controller is not None:
                command_controller.update()
                obs = env.get_observations()
            # agent stepping
            actions = policy(obs)
            # env stepping
            obs, _, _, _ = env.step(actions)
            step_count += 1

    if args_cli.max_steps is not None:
        print(f"[INFO] Replay completed after {step_count} policy steps.")

    # close the simulator
    env.close()


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