from ament_index_python.packages import get_package_share_directory

from launch import LaunchDescription
from launch.actions import DeclareLaunchArgument, GroupAction, OpaqueFunction
from launch.substitutions import LaunchConfiguration, PathJoinSubstitution
from launch_ros.actions import Node
from launch_ros.descriptions import ParameterFile
from launch.substitutions import PythonExpression
from nav2_common.launch import RewrittenYaml
import os


ARGUMENTS = [
    DeclareLaunchArgument('sim_mode', default_value='false'),
    DeclareLaunchArgument('robot', default_value='fr09', choices=['mkmini', 'fr09']),
    DeclareLaunchArgument('database_path', default_value=''),
]

def launch_setup(context, *args, **kwargs):
    robot = LaunchConfiguration('robot')
    sim_mode = LaunchConfiguration('sim_mode')
    database_path = LaunchConfiguration('database_path').perform(context).strip()
    if not database_path:
        raise RuntimeError('3D localization requires database_path:=/modules/maps/<saved-map>.db')
    if not os.path.isfile(database_path):
        raise RuntimeError(f'3D localization database was not found: {database_path}')

    remappings = [
        ('/tf', 'tf'),
        ('/tf_static', 'tf_static'),
        ('odom', 'icp_odom'),
        ('scan_cloud', '/lidar_filtered'),
    ]
    # Nav2 map_server owns /map during navigation and publishes the complete
    # saved PGM map. Keep RTAB-Map's local occupancy update separate so it
    # cannot race map_server and replace the navigation map intermittently.
    rtabmap_remappings = [*remappings, ('map', '/rtabmap/map')]
    
    # Directories
    pkg_localization = get_package_share_directory(
        'localization_pkg')
    
    rtabmap_params_file = PathJoinSubstitution([
        pkg_localization,
        'config',
        PythonExpression(["'rtabmap_' + '", robot, "' + '.yaml'"])])
    param_substitutions = {'use_sim_time': sim_mode}
    configured_params = ParameterFile(
        RewrittenYaml(
            source_file=rtabmap_params_file,
            param_rewrites=param_substitutions,
            convert_types=True),
        allow_substs=True)
    
    load_nodes = GroupAction(
        actions=[
            Node(
                package='rtabmap_odom',
                executable='icp_odometry',
                name='icp_odometry',
                output='screen',
                parameters=[configured_params, {'use_sim_time': sim_mode}],
                remappings=remappings,
            ),
            Node(
                package='rtabmap_slam',
                executable='rtabmap',
                name='rtabmap',
                output='screen',
                parameters=[configured_params, {
                    'use_sim_time': sim_mode,
                    'database_path': database_path,
                    'Mem/IncrementalMemory': 'false',
                    'Mem/InitWMWithAllNodes': 'true',
                    'map_frame_id': 'map',
                    'odom_frame_id': 'icp_odom',
                    'publish_tf': True,
                }],
                remappings=rtabmap_remappings,
            ),
        ],
    )
    return [load_nodes]

def generate_launch_description():
    # Create launch description and add actions
    ld = LaunchDescription(ARGUMENTS)
    ld.add_action(OpaqueFunction(function=launch_setup))
    
    return ld