from artefacts_toolkit.config import get_artefacts_param

get_artefacts_param(
    param_type,
    param_name,
    default=None,
    is_ros=True
)

scenarios:
  defaults:
    output_dirs: ["output"]
    metrics:
        - /odometry_error
    params:
      launch/world: ["bookstore.sdf", "empty.sdf"]

def generate_test_description():
    try:
        world = get_artefacts_param("launch", "world")
    except FileNotFoundError:
        world = "empty.world"

    run_headless = LaunchConfiguration("run_headless")
    launch_navigation_stack = IncludeLaunchDescription(
        PythonLaunchDescriptionSource(
            [
                os.path.join(
                    get_package_share_directory("sam_bot_nav2_gz"),
                    "launch",
                    "complete_navigation.launch.py"
                ),
            ]
        ),
        launch_arguments=[("run_headless", run_headless), ("world_file", world)],
    )

...# Rest of launch test file   

from artefacts_toolkit.chart import make_chart

make_chart(
    filepath,
    topic_x,
    topic_y,
    field_unit=None,
    output_dir="output",
    chart_name="chart",
    file_type="rosbag"
)

# my_test_file.launch.py
# test code

...

@launch_testing.post_shutdown_test()
class TestProcOutputAfterShutdown(unittest.TestCase):
    def test_exit_code(self, rosbag_filepath):
        make_chart(
            rosbag_filepath,
            "/odom.pose.pose.position.x",
            "/odom.pose.pose.position.y",
            field_unit="m",
            chart_name="odometry_position",
        )

from artefacts_toolkit.gazebo import bridge, gz

bridge.get_camera_bridge(
    topic_name,
    condition=None
)

@pytest.mark.launch_test
def generate_test_description():
    from artefacts_toolkit.gazebo import bridge
    ...

    camera_topic = "/observation_camera/image"
    bag_recorder, rosbag_filepath = get_bag_recorder([camera_topic])
    sim = IncludeLaunchDescription(
        PythonLaunchDescriptionSource(["bringup", "/robot.launch.py"])
    )

    record_video_launch_arg = DeclareLaunchArgument(
        "record_video", default_value="true"
    )
    record_video = LaunchConfiguration("record_video")
    camera_bridge = bridge.get_camera_bridge(camera_topic, condition=IfCondition(record_video))
    
    return LaunchDescription(
        [
            record_video_launch_arg,
            sim,
            camera_bridge,
            controller_process,
            launch_testing.actions.ReadyToTest(),
        ]
    )

gz.get_sim_objects(world_file)

objects = [
    {
        "name": "model_1",
        "pose": "0 0 0 0 0 0"
    }, 
    {
        "name": "model_2",
        "pose": "1 2 3 0 0 0"
    },
    ...
]

objects_positions = {
    "model_1": "0 0 0 0 0 0", 
    "model_2": "1 2 3 0 0 0",
    ...
}

# Working with the objects_position dict
_, objects_positions = gz.get_sim_objects("world.sdf")
model_pose = objects_positions["model_1"]  # "0 0 0 0 0 0"

# Working with the objects list
objects, _ = gz.get_sim_objects("world.sdf")
for model in objects:
    print(f"{model['name']} is at position {model['pose']}")

gz.get_model_location(model_name)

(x_position, y_position, z_position)

class TestProcOutput(unittest.TestCase):
    def test_moved_meatballs(self, proc_output, controller_process):
        # Original locations
        sim_objects, sim_objects_positions = gz.get_sim_objects("worlds/env.sdf") # this includes models poses
        meatball_models = [obj["name"] for obj in sim_objects if "karaage" in obj["name"]]
        #  Wait for the control loop to finish
        proc_output.assertWaitFor("Done with tasks execution", timeout=300)
        picked_meatballs = 0 # karaage moved for more than 10cm
        for meatball in meatball_models:
            x, y, z = gz.get_model_location(meatball)
            x_original, y_original, z_original = sim_objects_positions[meatball]
            dist = ((x - x_original) ** 2 + (y - y_original) ** 2 + (z - z_original) ** 2) ** 0.5
            if dist > 0.1: #10cm
                # Likely to have been picked and moved somewhere else
                picked_meatballs += 1
        self.assertEqual(picked_meatballs, 4)

gz.kill_gazebo()

from artefacts_toolkit.gazebo import gz
...

@launch_testing.post_shutdown_test()
class TestProcOutputAfterShutdown(unittest.TestCase):
    def test_exit_code(self, proc_info, controller_process, rosbag_filepath):
        gz.kill_gazebo()
        ...

from artefacts_toolkit.rosbag import get_bag_recorder, image_topics, message_topics

rosbag.get_bag_recorder(
    topic_names,
    directory="rosbags",
    use_sim_time=False
)

@pytest.mark.launch_test
def generate_test_description():
    camera_topics = ["/depth_cam/rgb"]
    extra_topics = ["/odom", "/noisy_estimate"]

    bag_recorder, rosbag_filepath = rosbag.get_bag_recorder(
        camera_topics + extra_topics, use_sim_time=False
    )

    test_odometry_node = ExecuteProcess(
        cmd=[
            "python3",
            "src/test_odometry_node.py",
        ]
    )
    return LaunchDescription(
        [
            test_odometry_node,
            launch_testing.actions.ReadyToTest(),
            bag_recorder,
        ]
    ), {
        "test_odometry_node": test_odometry_node,
        "rosbag_filepath": rosbag_filepath,
    }


@launch_testing.post_shutdown_test()
class TestProcOutputAfterShutdown(unittest.TestCase):
    def test_end_position(self, rosbag_filepath):
        final_distance_from_start = message_topics.get_final_message(
            rosbag_filepath, "/distance_from_start.data"
        )

        assert final_distance_from_start < 0.1, (
            f"Final distance from start is {final_distance_from_start}, expected less than 10cm"
        )

image_topics.extract_camera_image(
    rosbag_file_path,
    camera_topic,
    output_dir="output"
)

from artefacts_toolkit.rosbag import rosbag, image_topics
def test_exit_code(
    self, proc_info, test_odometry_node, rosbag_filepath
):
    ...

    image_topics.extract_camera_image(rosbag_filepath, "/depth_cam/rgb")

image_topics.extract_video(
    bag_path,
    topic_name,
    output_path,
    frame_rate=20
)

from artefacts_toolkit.rosbag import rosbag, image_topics
def test_exit_code(
    self, proc_info, test_odometry_node, rosbag_filepath
):
    ...

    image_topics.extract_video(rosbag_filepath, "/depth_cam/rgb", "output/depth_cam.webm")

message_topics.get_final_message(
    rosbag_filepath,
    topic
)

@launch_testing.post_shutdown_test()
class TestProcOutputAfterShutdown(unittest.TestCase):
    def test_end_position(self, rosbag_filepath):
        final_distance_from_start = message_topics.get_final_message(
            rosbag_filepath, "/distance_from_start.data"
        )

        assert final_distance_from_start < 0.1, (
            f"Final distance from start is {final_distance_from_start}, expected less than 10cm"
        )