# Requirements:
#   Install Turtlebot3 packages
#   Modify turtlebot3_waffle SDF:
#     1) Edit /opt/ros/$ROS_DISTRO/share/turtlebot3_gazebo/models/turtlebot3_waffle/model.sdf
#     2) We can increase min scan range from 0.12 to 0.2 to avoid having scans 
#        hitting the robot itself
#
# Example:
#   $ ros2 launch rtabmap_demos turtlebot3_sim_scan_demo.launch.py
#
#   Teleop:
#     $ ros2 run turtlebot3_teleop teleop_keyboard

from ament_index_python.packages import get_package_share_directory

from launch import LaunchDescription
from launch.actions import AppendEnvironmentVariable, DeclareLaunchArgument, IncludeLaunchDescription, OpaqueFunction
from launch.substitutions import LaunchConfiguration, PathJoinSubstitution
from launch.launch_description_sources import PythonLaunchDescriptionSource
from launch_ros.substitutions import FindPackageShare

import os

ROS_DISTRO = os.environ.get('ROS_DISTRO')

def launch_setup(context, *args, **kwargs):
    if not 'TURTLEBOT3_MODEL' in os.environ:
        os.environ['TURTLEBOT3_MODEL'] = 'waffle'

    # Directories
    pkg_turtlebot3_gazebo = get_package_share_directory(
        'turtlebot3_gazebo')
    pkg_nav2_bringup = get_package_share_directory(
        'nav2_bringup')
    pkg_rtabmap_demos = get_package_share_directory(
        'rtabmap_demos')
    
    world_name = LaunchConfiguration('world').perform(context)
    
    icp_odometry = LaunchConfiguration('icp_odometry').perform(context)
    icp_odometry = icp_odometry == 'True' or icp_odometry == 'true'
    if icp_odometry:
        # modified nav2 params to use icp_odom instead odom frame
        if ROS_DISTRO == 'humble':
            nav2_params_file = PathJoinSubstitution(
                [FindPackageShare('rtabmap_demos'), 'params', 'humble', 'turtlebot3_scan_nav2_params.yaml']
            )
        else:
            nav2_params_file = PathJoinSubstitution(
                [FindPackageShare('rtabmap_demos'), 'params', 'turtlebot3_scan_nav2_params.yaml']
            )
    else:
        if ROS_DISTRO == 'humble':
            # original nav2 params
            nav2_params_file = PathJoinSubstitution(
                [FindPackageShare('nav2_bringup'), 'params', 'nav2_params.yaml']
            )
        else:
            # original nav2 params but with "enable_stamped_cmd_vel: True"
            nav2_params_file = PathJoinSubstitution(
                [FindPackageShare('rtabmap_demos'), 'params', 'turtlebot3_nav2_params.yaml']
            )
            

    # Paths
    nav2_launch = PathJoinSubstitution(
        [pkg_nav2_bringup, 'launch', 'navigation_launch.py'])
    rviz_launch = PathJoinSubstitution(
        [pkg_nav2_bringup, 'launch', 'rviz_launch.py'])
    rtabmap_launch = PathJoinSubstitution(
        [pkg_rtabmap_demos, 'launch', 'turtlebot3', 'turtlebot3_scan.launch.py'])
    
    nav2 = IncludeLaunchDescription(
        PythonLaunchDescriptionSource([nav2_launch]),
        launch_arguments=[
            ('use_sim_time', 'true'),
            ('params_file', nav2_params_file)
        ]
    )
    rviz = IncludeLaunchDescription(
        PythonLaunchDescriptionSource([rviz_launch])
    )
    rtabmap = IncludeLaunchDescription(
        PythonLaunchDescriptionSource([rtabmap_launch]),
        launch_arguments=[
            ('localization', LaunchConfiguration('localization')),
            ('use_sim_time', 'true')
        ]
    )
        
    # To use ICP odometry, we should increase clock rate of gazebo (humble), we copied content of
    # turtlebot3_gazebo/launch/turtlebot3_world.launch here.
    turtlebot3_nodes = []
    if ROS_DISTRO == 'humble':
        pkg_gazebo_ros = get_package_share_directory('gazebo_ros')

        import tempfile
        with tempfile.NamedTemporaryFile(mode='w+t', delete=False) as clock_override_file:
            clock_override_file.write("---\n"+
                    "gazebo:\n"+
                    "    ros__parameters:\n"+
                    "        publish_rate: 100.0")

        world = os.path.join(
                pkg_turtlebot3_gazebo,
                'worlds',
                f'turtlebot3_{world_name}.world'
        )

        gzserver_cmd = IncludeLaunchDescription(
            PythonLaunchDescriptionSource(
                os.path.join(pkg_gazebo_ros, 'launch', 'gzserver.launch.py')
            ),
            launch_arguments={
                'world': world,
                'params_file': clock_override_file.name}.items()
        )

        gzclient_cmd = IncludeLaunchDescription(
            PythonLaunchDescriptionSource(
                os.path.join(pkg_gazebo_ros, 'launch', 'gzclient.launch.py')
            )
        )
        
        robot_state_publisher_cmd = IncludeLaunchDescription(
            PythonLaunchDescriptionSource(
                os.path.join(pkg_turtlebot3_gazebo, 'launch', 'robot_state_publisher.launch.py')
            ),
            launch_arguments={'use_sim_time': 'true'}.items()
        )
        spawn_turtlebot_cmd = IncludeLaunchDescription(
            PythonLaunchDescriptionSource(
                os.path.join(pkg_turtlebot3_gazebo, 'launch', 'spawn_turtlebot3.launch.py')
            ),
            launch_arguments={
                'x_pose': LaunchConfiguration('x_pose'),
                'y_pose': LaunchConfiguration('y_pose')
            }.items()
        )
        
        set_env_vars_resources = AppendEnvironmentVariable(
                'GZ_SIM_RESOURCE_PATH',
                os.path.join(pkg_turtlebot3_gazebo, 'models'))
        turtlebot3_nodes = [
            gzserver_cmd,
            gzclient_cmd,
            robot_state_publisher_cmd,
            spawn_turtlebot_cmd,
            set_env_vars_resources
        ]
    else:
        gazebo_launch = PathJoinSubstitution([pkg_turtlebot3_gazebo, 'launch', f'turtlebot3_{world_name}.launch.py'])
        gazebo = IncludeLaunchDescription(
            PythonLaunchDescriptionSource([gazebo_launch]),
            launch_arguments=[
                ('x_pose', LaunchConfiguration('x_pose')),
                ('y_pose', LaunchConfiguration('y_pose'))
            ]
        )
        turtlebot3_nodes = [gazebo]

    return [
        # Nodes to launch
        nav2,
        rviz,
        rtabmap] + turtlebot3_nodes

def generate_launch_description():
    return LaunchDescription([
        
        # Launch arguments
        DeclareLaunchArgument(
            'localization', default_value='false',
            description='Launch in localization mode.'),
        
        DeclareLaunchArgument(
            'world', default_value='world',
            choices=['world', 'house', 'dqn_stage1', 'dqn_stage2', 'dqn_stage3', 'dqn_stage4'],
            description='Turtlebot3 gazebo world.'),
        
        DeclareLaunchArgument(
            'icp_odometry', default_value='false',
            description='Launch ICP odometry on top of wheel odometry.'),
        
        DeclareLaunchArgument(
            'deskewing', default_value='false',
            description='Do lidar scan deskewing based on wheel odometry. In simulation it '
                        'doesn\'t matter because laser scans are not skewed, but the option '
                        'is there for the demo.'),

        DeclareLaunchArgument(
            'x_pose', default_value='-2.0',
            description='Initial position of the robot in the simulator.'),
        
        DeclareLaunchArgument(
            'y_pose', default_value='0.5',
            description='Initial position of the robot in the simulator.'),

        OpaqueFunction(function=launch_setup)
    ])