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   

get_artefacts_params()

scenarios:
    defaults:
        params:
            headless: "True"

@launch_pytest.fixture(scope="module")
def launch_description(rosbag_recording):
    pkg_path = Path(get_package_share_directory("my_package"))
    try:
        headless = get_artefacts_params().get("headless", "False")
    except RuntimeError:
        # If artefacts params are not available, default to False
        headless = "False"

    start_launch = IncludeLaunchDescription(
        PythonLaunchDescriptionSource(
            str(pkg_path / "launch" / "mytestfile.launch.py")
        ),
        launch_arguments={
            "headless": headless,
        }.items(),
    )

merge_ros_params_files(
    source,
    override,
    destination,
    rosify=False
)

# source
outer:
    inner1: 1
    inner2: 
        leaf: keep
flat: src

# override
outer:
    inner2:
        leaf: replaced
        new_leaf: 321
    inner3: added
flat: override

# rosify=False (Default)
outer:
    inner1: 1
    inner2:
        leaf: replaced
        new_leaf: 321
    inner3: added
flat: override

#source
controller_server:
    ros__parameters:
        enable_stamped_cmd_vel: true
        controller_frequency: 20.0
        min_x_velocity_threshold: 0.001

#override
controller_server/controller_frequency: 30
controller_server/min_x_velocity_threshold: 0.005
controller_server/min_y_velocity_threshold: 0.5   
    
#destination
controller_server:
    ros__parameters:
        enable_stamped_cmd_vel: true
        controller_frequency: 30
        min_x_velocity_threshold: 0.005
        min_y_velocity_threshold: 0.5

from artefacts_toolkit.chart import make_chart

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

# 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",
            output_format="csv"
        )

make_chart(
    rosbag_filepath,
    "time",
    ["/critics_cost/ConstraintCritic.data", "/critics_cost/GoalAngleCritic.data", "/critics_cost/GoalCritic.data", "/critics_cost/CostCritic.data", "/critics_cost/PathAlignCritic.data", "/critics_cost/PathAngleCritic.data", "/critics_cost/PathFollowCritic.data", "/critics_cost/PreferForwardCritic.data"],
    axis_x_name="time",
    axis_y_name="critics",
    chart_name="Critics vs Time",
)

 make_chart(
        rosbag_filepath,
        topic_x=["/odom.pose.pose.position.x", "/gt_odom.pose.pose.position.x"],
        topic_y=["/odom.pose.pose.position.y", "/gt_odom.pose.pose.position.y"],
        axis_x_name="x (m)",
        axis_y_name="y (m)",
        chart_name="odom vs gt positions"
    )

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.add_entity(world_xml, entity_xml)

from artefacts_toolkit.gazebo import gz

# Load existing world file
world_xml = gz.load_world("worlds/env.sdf")

# Create an actor with waypoints
waypoints = [
    [4, -3.2, 1.0],
    [0.8, -3.2, 1.0],
    [0.8, -6, 1.0],
    [4, -6, 1.0],
]
actor_xml = gz.make_actor("pedestrian_1", waypoints)

# Add the actor to the world
gz.add_entity(world_xml, actor_xml)

gz.load_world(world_file)

from artefacts_toolkit.gazebo import gz

# Load a world file
world_xml = gz.load_world("worlds/env.sdf")

# Now you can modify the world and add entities
actor_xml = gz.make_actor("pedestrian", [[0, 0, 0], [1, 1, 0]])
gz.add_entity(world_xml, actor_xml)

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()
        ...

gz.make_actor(
    actor_name,
    waypoints,
    walk_speed=0.8,
    rotate_speed=1.8,
    enable_loop=True,
    output_type="xml"
)

<actor name="{actor_name}">
    <skin>
        <filename>https://fuel.gazebosim.org/1.0/Mingfei/models/actor/tip/files/meshes/walk.dae</filename>
    </skin>
    
    <animation name="walking">
        <filename>https://fuel.gazebosim.org/1.0/Mingfei/models/actor/tip/files/meshes/walk.dae</filename>
        <interpolate_x>true</interpolate_x>
    </animation>

    <script>
        <loop>true / false</loop>
        <delay_start>0.0</delay_start>
        <auto_start>true</auto_start>
        
        <trajectory id="0" type="walking">
            <!-- The generated waypoints given the provided inputs -->
        </trajectory>
    </script>
</actor>

waypoints = [
    [4, -3.2, 1.0],
    [0.8, -3.2, 1.0],
    [0.8, -6, 1.0],
    [4, -6, 1.0],
]

actor_xml = gz.make_actor("my_actor", waypoints, walk_speed=0.8, rotate_speed=1.8, enable_loop=True)

<actor name="my_actor">
    <skin>
        <filename>https://fuel.gazebosim.org/1.0/Mingfei/models/actor/tip/files/meshes/walk.dae</filename>
    </skin>
    
    <animation name="walking">
        <filename>https://fuel.gazebosim.org/1.0/Mingfei/models/actor/tip/files/meshes/walk.dae</filename>
        <interpolate_x>true</interpolate_x>
    </animation>

    <script>
        <loop>true</loop>
        <delay_start>0.0</delay_start>
        <auto_start>true</auto_start>
        
        <trajectory id="0" type="walking">
            <waypoint>
                <time>0</time>
                <pose>4 -3.2 1.0 0 0 3.141592653589793</pose>
            </waypoint>
            <waypoint>
                <time>4.0</time>
                <pose>0.8 -3.2 1.0 0 0 3.141592653589793</pose>
            </waypoint>
            <waypoint>
                <time>6.617993877991495</time>
                <pose>0.8 -3.2 1.0 0 0 -1.5707963267948966</pose>
            </waypoint>
            <waypoint>
                <time>10.117993877991495</time>
                <pose>0.8 -6 1.0 0 0 -1.5707963267948966</pose>
            </waypoint>
            <waypoint>
                <time>10.99065850398866</time>
                <pose>0.8 -6 1.0 0 0 0.0</pose>
            </waypoint>
            <waypoint>
                <time>14.99065850398866</time>
                <pose>4 -6 1.0 0 0 0.0</pose>
            </waypoint>
            <waypoint>
                <time>15.863323129985826</time>
                <pose>4 -6 1.0 0 0 1.5707963267948966</pose>
            </waypoint>
            <waypoint>
                <time>19.363323129985826</time>
                <pose>4 -3.2 1.0 0 0 1.5707963267948966</pose>
            </waypoint>
            <waypoint>
                <time>20.23598775598299</time>
                <pose>4 -3.2 1.0 0 0 3.141592653589793</pose>
            </waypoint>
            
        </trajectory>
    </script>
</actor>

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"
        )

from artefacts_toolkit.rerun import recorder, reader, video

recorder.get_output_dir(
    directory=None
)

from artefacts_toolkit.rerun import recorder

OUTPUT_FOLDER = recorder.get_output_dir()
rrd_path = OUTPUT_FOLDER / "my-rerun-recording.rrd"

recorder.start_recording(
    application_id,
    filename="recording.rrd",
    directory=None,
    handle_sigterm=True
)

import rerun as rr
import simulator as sim
from artefacts_toolkit.rerun import recorder


class RerunSimulator(sim.SimulatorMujoco):
    def _set_sim_time(self):
        rr.set_time("sim_time", duration=float(self.mujoco_data.time))

    def _log_joy(self, joy):
        self._set_sim_time()
        rr.log("cmd/fwd", rr.Scalars(joy.axes[1]))
        rr.log("cmd/yaw", rr.Scalars(joy.axes[2]))

    def _read_state(self):
        super()._read_state()
        q = self.mujoco_data.qpos  # [0:3] base xyz, [3:7] base quat wxyz
        w, x, y, z = q[3:7]
        self._set_sim_time()
        rr.log("base", rr.Transform3D(translation=q[0:3],
                                      quaternion=rr.Quaternion(xyzw=[x, y, z, w])))
        rr.log("base/yaw", rr.Scalars(yaw_from_quat(w, x, y, z)))


def main():
    robot = sim.Robot(sim.RobotType.Tron2, True)
    robot.init("127.0.0.1")

    # After robot.init(): the SDK resets signal handlers
    recorder.start_recording("tron2_loop", "recording-loop.rrd")

    RerunSimulator(...).run()

reader.load(
    rrd
)

from artefacts_toolkit.rerun import reader


@pytest.fixture(scope="module")
def recording(recording_path):
    assert recording_path.exists(), f"Recording not found at {recording_path}"
    return reader.load(recording_path)

reader.get_entity_paths(
    rrd
)

assert "/bodies/uwb_tag" in reader.get_entity_paths(recording), "target was never logged"

reader.get_timelines(
    rrd
)

reader.get_columns(
    rrd,
    entity=None
)

>>> reader.get_columns(recording, "/base")
['/base:Transform3D:quaternion', '/base:Transform3D:translation']

reader.get_column(
    rrd,
    column,
    timeline=None
)

times, yaw = reader.get_column(recording, "/base/yaw:Scalars:scalars")
assert times[-1] > 60.0, "mission ended early"
assert np.abs(yaw).max() <= np.pi

reader.get_final_message(
    rrd,
    column
)

wp_col = "/network/mission_waypoints:Points3D:positions"
assert wp_col in reader.get_columns(recording), f"Could not find {wp_col} in recording"

track = np.vstack(reader.get_final_message(recording, wp_col))[:, :2]  # N x 2

reader.get_message_count(
    rrd,
    entity
)

assert reader.get_message_count(recording, "/cam") >= 100, "camera stopped publishing"

reader.to_dataframe(
    rrd,
    contents=None,
    index=None,
    step=None,
    fill_latest_at=True
)

@pytest.fixture(scope="module")
def df(recording_path):
    """Whole recording resampled onto a uniform 20 ms sim-time grid."""
    frame = reader.to_dataframe(recording_path, ["/base/**", "/cmd/**"], index="sim_time", step=0.02)

    out = frame.dropna(subset=["/base:Transform3D:translation"]).reset_index(drop=True)
    out["x"] = [t[0] for t in out["/base:Transform3D:translation"]]
    out["y"] = [t[1] for t in out["/base:Transform3D:translation"]]
    out["yaw"] = out["/base/yaw:Scalars:scalars"]
    out["cmd_fwd"] = out["/cmd/fwd:Scalars:scalars"].fillna(0.0)
    return out


def test_straights_cover_side_length(df, straights):
    for i, (a, b) in enumerate(straights):
        dist = math.hypot(df["x"][b - 1] - df["x"][a], df["y"][b - 1] - df["y"][a])
        assert abs(dist - SIDE_LENGTH) < DIST_TOLERANCE

video.extract_video(
    rrd,
    entity,
    output_path,
    timeline=None,
    frame_rate=20
)

from artefacts_toolkit.rerun import recorder, video

video.extract_video(recording, "/cam", recorder.get_output_dir() / "head_camera.mp4")

video.extract_camera_image(
    rrd,
    entity,
    output_dir="output"
)

video.extract_camera_image(recording, "/cam", recorder.get_output_dir())