import os
from os import cpu_count

import xacro
from ament_index_python.packages import get_package_share_directory
from launch import LaunchDescription
from launch.actions import (
    DeclareLaunchArgument,
    IncludeLaunchDescription,
    OpaqueFunction,
    SetEnvironmentVariable,
    TimerAction,
)
from launch.launch_description_sources import PythonLaunchDescriptionSource
from launch.substitutions import LaunchConfiguration, PathJoinSubstitution
from launch_ros.actions import Node
from launch_ros.substitutions import FindPackageShare


def launch_setup(context, *args, **kwargs):
    # Initialize Arguments
    robot_type = LaunchConfiguration("robot_type")
    # General arguments
    with_gripper_condition = LaunchConfiguration("with_gripper")
    controllers_file = LaunchConfiguration("controllers_file")
    cartesian_controllers = LaunchConfiguration("cartesian_controllers")
    description_package = LaunchConfiguration("description_package")
    description_file = LaunchConfiguration("description_file")
    robot_name = LaunchConfiguration("robot_name")
    start_joint_controller = LaunchConfiguration("start_joint_controller")
    initial_joint_controller = LaunchConfiguration("initial_joint_controller")
    launch_simulation = LaunchConfiguration("launch_sim")
    launch_moveit = LaunchConfiguration("launch_moveit")
    launch_task_planner = LaunchConfiguration("launch_task_planner")
    launch_perception = LaunchConfiguration("launch_perception")
    moveit_config_package = LaunchConfiguration("moveit_config_package")
    moveit_config_file = LaunchConfiguration("moveit_config_file")
    use_sim_time = LaunchConfiguration("use_sim_time")
    sim_gazebo = LaunchConfiguration("sim_gazebo")
    hardware = LaunchConfiguration("hardware")
    env_manager = LaunchConfiguration("env_manager")
    launch_controllers = LaunchConfiguration("launch_controllers")
    gripper_name = LaunchConfiguration("gripper_name")

    # training arguments
    env = LaunchConfiguration("env")
    use_sim_time = LaunchConfiguration("use_sim_time")
    log_level = LaunchConfiguration("log_level")
    env_kwargs = LaunchConfiguration("env_kwargs")

    sim_gazebo = LaunchConfiguration("sim_gazebo")
    launch_simulation = LaunchConfiguration("launch_sim")

    initial_joint_controllers_file_path = os.path.join(
        get_package_share_directory("rbs_arm"), "config", "controllers.yaml"
    )

    xacro_file = os.path.join(
        get_package_share_directory(description_package.perform(context)),
        "urdf",
        description_file.perform(context),
    )

    robot_description_doc = xacro.process_file(
        xacro_file,
        mappings={
            "gripper_name": gripper_name.perform(context),
            "hardware": hardware.perform(context),
            "simulation_controllers": initial_joint_controllers_file_path,
            "namespace": "",
        },
    )

    robot_description_content = robot_description_doc.toprettyxml(indent="  ")
    robot_description = {"robot_description": robot_description_content}

    single_robot_setup = IncludeLaunchDescription(
        PythonLaunchDescriptionSource(
            [
                PathJoinSubstitution(
                    [FindPackageShare("rbs_bringup"), "launch", "rbs_robot.launch.py"]
                )
            ]
        ),
        launch_arguments={
            "env_manager": env_manager,
            "with_gripper": with_gripper_condition,
            "gripper_name": gripper_name,
            "controllers_file": initial_joint_controllers_file_path,
            "robot_type": robot_type,
            "cartesian_controllers": cartesian_controllers,
            "description_package": description_package,
            "description_file": description_file,
            "robot_name": robot_type,
            "start_joint_controller": start_joint_controller,
            "initial_joint_controller": initial_joint_controller,
            "launch_simulation": launch_simulation,
            "launch_moveit": launch_moveit,
            "launch_task_planner": launch_task_planner,
            "launch_perception": launch_perception,
            "moveit_config_package": moveit_config_package,
            "moveit_config_file": moveit_config_file,
            "use_sim_time": use_sim_time,
            "sim_gazebo": sim_gazebo,
            "hardware": hardware,
            "launch_controllers": "true",
            "robot_description": robot_description_content,
        }.items(),
    )

    args = [
        "--env",
        env,
        "--env-kwargs",
        env_kwargs,
        "--ros-args",
        "--log-level",
        log_level,
    ]

    rl_task = Node(
        package="rbs_gym",
        executable="test_agent",
        output="log",
        arguments=args,
        parameters=[{"use_sim_time": True}, robot_description],
    )

    clock_bridge = Node(
        package="ros_gz_bridge",
        executable="parameter_bridge",
        arguments=["/clock@rosgraph_msgs/msg/Clock[ignition.msgs.Clock"],
        output="screen",
    )

    delay_robot_control_stack = TimerAction(period=10.0, actions=[single_robot_setup])

    nodes_to_start = [
        # env,
        rl_task,
        clock_bridge,
        delay_robot_control_stack,
    ]
    return nodes_to_start


def generate_launch_description():
    declared_arguments = []
    declared_arguments.append(
        DeclareLaunchArgument(
            "robot_type",
            description="Type of robot by name",
            choices=["rbs_arm", "ur3", "ur3e", "ur5", "ur5e", "ur10", "ur10e", "ur16e"],
            default_value="rbs_arm",
        )
    )
    # General arguments
    declared_arguments.append(
        DeclareLaunchArgument(
            "controllers_file",
            default_value="controllers_gazebosim.yaml",
            description="YAML file with the controllers configuration.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "description_package",
            default_value="rbs_arm",
            description="Description package with robot URDF/XACRO files. Usually the argument \
        is not set, it enables use of a custom description.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "description_file",
            default_value="rbs_arm_modular.xacro",
            description="URDF/XACRO description file with the robot.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "robot_name",
            default_value="arm0",
            description="Name for robot, used to apply namespace for specific robot in multirobot setup",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "start_joint_controller",
            default_value="false",
            description="Enable headless mode for robot control",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "initial_joint_controller",
            default_value="joint_trajectory_controller",
            description="Robot controller to start.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "moveit_config_package",
            default_value="rbs_arm",
            description="MoveIt config package with robot SRDF/XACRO files. Usually the argument \
        is not set, it enables use of a custom moveit config.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "moveit_config_file",
            default_value="rbs_arm.srdf.xacro",
            description="MoveIt SRDF/XACRO description file with the robot.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "use_sim_time",
            default_value="true",
            description="Make MoveIt to use simulation time.\
            This is needed for the trajectory planing in simulation.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "gripper_name",
            default_value="rbs_gripper",
            choices=["rbs_gripper", ""],
            description="choose gripper by name (leave empty if hasn't)",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "with_gripper", default_value="true", description="With gripper or not?"
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "sim_gazebo", default_value="true", description="Gazebo Simulation"
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "env_manager", default_value="false", description="Launch env_manager?"
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "launch_sim",
            default_value="true",
            description="Launch simulator (Gazebo)?\
                              Most general arg",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "launch_moveit", default_value="false", description="Launch moveit?"
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "launch_perception", default_value="false", description="Launch perception?"
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "launch_task_planner",
            default_value="false",
            description="Launch task_planner?",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "cartesian_controllers",
            default_value="true",
            description="Load cartesian\
                              controllers?",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "hardware",
            choices=["gazebo", "mock"],
            default_value="gazebo",
            description="Choose your harware_interface",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "launch_controllers",
            default_value="true",
            description="Launch controllers?",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "gazebo_gui", default_value="true", description="Launch gazebo with gui?"
        )
    )
    # training arguments
    declared_arguments.append(
        DeclareLaunchArgument(
            "env",
            default_value="Reach-Gazebo-v0",
            description="Environment ID",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "env_kwargs",
            default_value="",
            description="Optional keyword argument to pass to the env constructor.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "vec_env",
            default_value="dummy",
            description="Type of VecEnv to use (dummy or subproc).",
        )
    )
    # Algorithm and training
    declared_arguments.append(
        DeclareLaunchArgument(
            "algo",
            default_value="sac",
            description="RL algorithm to use during the training.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "n_timesteps",
            default_value="-1",
            description="Overwrite the number of timesteps.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "hyperparams",
            default_value="",
            description="Optional RL hyperparameter overwrite (e.g. learning_rate:0.01 train_freq:10).",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "num_threads",
            default_value="-1",
            description="Number of threads for PyTorch (-1 to use default).",
        )
    )
    # Continue training an already trained agent
    declared_arguments.append(
        DeclareLaunchArgument(
            "trained_agent",
            default_value="",
            description="Path to a pretrained agent to continue training.",
        )
    )
    # Random seed
    declared_arguments.append(
        DeclareLaunchArgument(
            "seed",
            default_value="-1",
            description="Random generator seed.",
        )
    )
    # Saving of model
    declared_arguments.append(
        DeclareLaunchArgument(
            "save_freq",
            default_value="10000",
            description="Save the model every n steps (if negative, no checkpoint).",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "save_replay_buffer",
            default_value="False",
            description="Save the replay buffer too (when applicable).",
        )
    )
    # Pre-load a replay buffer and start training on it
    declared_arguments.append(
        DeclareLaunchArgument(
            "preload_replay_buffer",
            default_value="",
            description="Path to a replay buffer that should be preloaded before starting the training process.",
        )
    )
    # Logging
    declared_arguments.append(
        DeclareLaunchArgument(
            "log_folder",
            default_value="logs",
            description="Path to the log directory.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "tensorboard_log",
            default_value="tensorboard_logs",
            description="Tensorboard log dir.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "log_interval",
            default_value="-1",
            description="Override log interval (default: -1, no change).",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "uuid",
            default_value="False",
            description="Ensure that the run has a unique ID.",
        )
    )
    # Evaluation
    declared_arguments.append(
        DeclareLaunchArgument(
            "eval_freq",
            default_value="-1",
            description="Evaluate the agent every n steps (if negative, no evaluation).",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "eval_episodes",
            default_value="5",
            description="Number of episodes to use for evaluation.",
        )
    )
    # Verbosity
    declared_arguments.append(
        DeclareLaunchArgument(
            "verbose",
            default_value="1",
            description="Verbose mode (0: no output, 1: INFO).",
        )
    )
    # HER specifics
    declared_arguments.append(
        DeclareLaunchArgument(
            "truncate_last_trajectory",
            default_value="True",
            description="When using HER with online sampling the last trajectory in the replay buffer will be truncated after) reloading the replay buffer.",
        )
    )
    declared_arguments.append(
        DeclareLaunchArgument(
            "log_level",
            default_value="error",
            description="The level of logging that is applied to all ROS 2 nodes launched by this script.",
        )
    )
    cpu_c = cpu_count()
    if cpu_c is not None:
        env_variables = [
            SetEnvironmentVariable(name="OMP_DYNAMIC", value="TRUE"),
            SetEnvironmentVariable(name="OMP_NUM_THREADS", value=str(cpu_c // 2)),
        ]
    else:
        env_variables = [
            SetEnvironmentVariable(name="OMP_DYNAMIC", value="TRUE"),
        ]


    return LaunchDescription(
        declared_arguments + [OpaqueFunction(function=launch_setup)] + env_variables
    )