#!/usr/bin/env python3
"""Exercise the real robot_localization EKF using a supplied straight-line fixture."""

import argparse
import csv
import json
import math
import os
from pathlib import Path
import signal
import subprocess
import tempfile
import time

import rclpy
from rclpy.node import Node
from nav_msgs.msg import Odometry
from sensor_msgs.msg import Imu
from rosgraph_msgs.msg import Clock
from tf2_msgs.msg import TFMessage
from rcl_interfaces.srv import GetParameters
import yaml

HERE = Path(__file__).resolve().parent
PACKAGE_VERSION = "3.8.3-1noble.20260916.090038"
NS = 1_000_000_000
START_NS = 100 * NS
STEP_NS = NS // 10
STEPS = 50
TIMEOUT = 15.0
CASES = ("baseline", "fuse-biased-gyro")
FIELDS = ("x", "y", "z", "roll", "pitch", "yaw", "vx", "vy", "vz",
          "vroll", "vpitch", "vyaw", "ax", "ay", "az")


def require(condition, description):
    if not condition:
        raise AssertionError(description)


def stamp_ns(stamp):
    return stamp.sec * NS + stamp.nanosec


def set_stamp(stamp, nanoseconds):
    stamp.sec, stamp.nanosec = divmod(nanoseconds, NS)


def yaw_of(q):
    return math.atan2(2 * (q.w * q.z + q.x * q.y),
                      1 - 2 * (q.y * q.y + q.z * q.z))


def covariance(size, diagonal):
    values = [0.0] * (size * size)
    for index, value in enumerate(diagonal):
        values[index * size + index] = value
    return values


def samples(index):
    timestamp = START_NS + index * STEP_NS
    odom = Odometry()
    set_stamp(odom.header.stamp, timestamp)
    odom.header.frame_id = "odom"
    odom.child_frame_id = "base_link"
    # This supplied pose is not selected by either filter. No pose or heading
    # measurement is fused: the estimator integrates the selected body twist.
    odom.pose.pose.position.x = index * 0.1
    odom.pose.pose.orientation.w = 1.0
    odom.pose.covariance = covariance(6, [1.0] * 6)
    odom.twist.twist.linear.x = 1.0
    odom.twist.covariance = covariance(6, [0.01] * 6)
    imu = Imu()
    set_stamp(imu.header.stamp, timestamp)
    imu.header.frame_id = "base_link"
    imu.orientation.w = 1.0
    imu.orientation_covariance[0] = -1.0
    imu.linear_acceleration_covariance[0] = -1.0
    imu.angular_velocity.z = 0.2
    imu.angular_velocity_covariance = covariance(3, [0.01, 0.01, 0.0001])
    return odom, imu


def supplied_row(index):
    return {
        "stamp_ns": START_NS + index * STEP_NS,
        "elapsed_s": index / 10,
        "reference_x_m": index / 10,
        "reference_y_m": 0.0,
        "reference_yaw_rad": 0.0,
        "wheel_vx_m_s": 1.0,
        "wheel_vy_m_s": 0.0,
        "wheel_yaw_rate_rad_s": 0.0,
        "wheel_selected_variance": 0.01,
        "imu_yaw_rate_rad_s": 0.2,
        "imu_yaw_rate_variance": 0.0001,
    }


def save_csv(path, rows):
    require(bool(rows), f"No rows to save: {path.name}")
    with path.open("w", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)


def spin_until(node, predicate, description, process=None, timeout=TIMEOUT):
    deadline = time.monotonic() + timeout
    while not predicate():
        require(time.monotonic() < deadline, f"Timed out: {description}")
        if process is not None:
            require(process.poll() is None, f"Estimator exited early: {process.returncode}")
        rclpy.spin_once(node, timeout_sec=0.01)


def settle(node, seconds):
    deadline = time.monotonic() + seconds
    while time.monotonic() < deadline:
        rclpy.spin_once(node, timeout_sec=0.01)


class Fixture(Node):
    def __init__(self):
        super().__init__("localization_lesson_fixture")
        self.current_ns = START_NS - STEP_NS
        self.rows = []
        self.transforms = []
        self.odom_pub = self.create_publisher(Odometry, "/fixture/odom", 20)
        self.imu_pub = self.create_publisher(Imu, "/fixture/imu", 20)
        self.clock_pub = self.create_publisher(Clock, "/clock", 10)
        self.odom_sub = self.create_subscription(
            Odometry, "/odometry/filtered", self.receive, 100)
        self.tf_sub = self.create_subscription(TFMessage, "/tf", self.receive_tf, 100)
        # This fixture keeps its system clock so discovery and cleanup can run
        # even before the estimator has received its first simulation clock.
        self.timer = self.create_timer(0.01, self.publish_clock)

    def publish_clock(self):
        message = Clock()
        set_stamp(message.clock, self.current_ns)
        self.clock_pub.publish(message)

    def receive(self, message):
        pose = message.pose.pose
        q = pose.orientation
        row = {
            "stamp_ns": stamp_ns(message.header.stamp),
            "elapsed_s": (stamp_ns(message.header.stamp) - START_NS) / NS,
            "frame": message.header.frame_id,
            "child_frame": message.child_frame_id,
            "x_m": pose.position.x, "y_m": pose.position.y,
            "z_m": pose.position.z, "yaw_rad": yaw_of(q),
            "qx": q.x, "qy": q.y, "qz": q.z, "qw": q.w,
            "vx_m_s": message.twist.twist.linear.x,
            "vy_m_s": message.twist.twist.linear.y,
            "yaw_rate_rad_s": message.twist.twist.angular.z,
            "x_variance": message.pose.covariance[0],
            "y_variance": message.pose.covariance[7],
            "yaw_variance": message.pose.covariance[35],
        }
        self.rows.append(row)

    def receive_tf(self, message):
        for transform in message.transforms:
            if transform.header.frame_id != "odom" or transform.child_frame_id != "base_link":
                continue
            translation = transform.transform.translation
            q = transform.transform.rotation
            self.transforms.append({
                "stamp_ns": stamp_ns(transform.header.stamp),
                "frame": transform.header.frame_id,
                "child_frame": transform.child_frame_id,
                "x_m": translation.x, "y_m": translation.y, "z_m": translation.z,
                "qx": q.x, "qy": q.y, "qz": q.z, "qw": q.w,
            })


def stop(process, stream):
    try:
        os.killpg(process.pid, signal.SIGINT)
    except ProcessLookupError:
        pass
    if process.poll() is None:
        try:
            process.wait(timeout=10)
        except subprocess.TimeoutExpired:
            os.killpg(process.pid, signal.SIGKILL)
            process.wait(timeout=5)
    stream.close()
    return process.returncode


def read_parameters(fixture, process, config, work, case):
    # Query only configured values. This version rejects a whole GetParameters
    # request if it includes an uninitialized optional covariance parameter.
    names = list(config)
    client = fixture.create_client(GetParameters, "/ekf_filter_node/get_parameters")
    spin_until(fixture, client.service_is_ready, "estimator parameter service", process)
    request = GetParameters.Request()
    request.names = names
    future = client.call_async(request)
    spin_until(fixture, future.done, "estimator parameter values", process)
    require(future.exception() is None, "Parameter readback failed")
    attributes = {1: "bool_value", 2: "integer_value", 3: "double_value",
                  4: "string_value", 6: "bool_array_value", 7: "integer_array_value",
                  8: "double_array_value", 9: "string_array_value"}
    actual = {}
    require(len(future.result().values) == len(names), "Parameter service omitted values")
    for name, value in zip(names, future.result().values):
        require(value.type in attributes, f"Unset or unsupported parameter: {name}")
        field = getattr(value, attributes[value.type])
        actual[name] = list(field) if value.type >= 6 else field
    for name, expected in config.items():
        require(actual[name] == expected, f"Loaded parameter differs from YAML: {name}")
    (work / f"{case}-parameters.json").write_text(json.dumps(actual, indent=2) + "\n")
    fixture.destroy_client(client)
    return actual


def run_case(case, work):
    fixture = Fixture()
    config = HERE / f"{case}.yaml"
    command = ["ros2", "run", "robot_localization", "ekf_node", "--ros-args",
               "-r", "__node:=ekf_filter_node", "--params-file", str(config)]
    print(f"$ ros2 run robot_localization ekf_node --ros-args "
          f"-r __node:=ekf_filter_node --params-file {case}.yaml", flush=True)
    stream = (work / f"{case}.log").open("w")
    process = subprocess.Popen(command, stdout=stream, stderr=subprocess.STDOUT,
                               start_new_session=True)
    completed = False
    try:
        spin_until(fixture, lambda: fixture.odom_pub.get_subscription_count() == 1
                   and fixture.count_publishers("/odometry/filtered") == 1
                   and (case == "baseline" or fixture.imu_pub.get_subscription_count() == 1),
                   "estimator discovery", process)
        settle(fixture, 0.5)
        params = yaml.safe_load(config.read_text())["ekf_filter_node"]["ros__parameters"]
        read_parameters(fixture, process, params, work, case)
        for index in range(STEPS + 1):
            odom, imu = samples(index)
            fixture.odom_pub.publish(odom)
            fixture.imu_pub.publish(imu)
            # Queue the next stamped measurements before advancing time. The
            # estimator's ROS-time timer needs an advancing clock to execute.
            settle(fixture, 0.05)
            fixture.current_ns = START_NS + index * STEP_NS
            fixture.publish_clock()
            # Hold this simulated time so independent sensor and clock callbacks
            # can run before advancing. This is not a transport timing guarantee.
            settle(fixture, 0.12)
            spin_until(fixture, lambda: any(row["stamp_ns"] >= fixture.current_ns
                                          for row in fixture.rows),
                       f"filter output for input {index}", process)
        settle(fixture, 0.3)
        final_ns = START_NS + STEPS * STEP_NS
        final_rows = [row for row in fixture.rows if row["stamp_ns"] == final_ns]
        require(bool(final_rows), "Missing final output stamp")
        final = final_rows[-1]
        for row in fixture.rows:
            require(row["frame"] == "odom" and row["child_frame"] == "base_link",
                    "Filtered odometry frames changed")
            require(all(math.isfinite(value) for value in row.values()
                        if isinstance(value, (int, float))), "Non-finite estimator output")
            require(abs(row["z_m"]) < 1e-8 and abs(row["qx"]) < 1e-8
                    and abs(row["qy"]) < 1e-8, "Planar state escaped two_d_mode")
            require(abs(sum(row[key] ** 2 for key in ("qx", "qy", "qz", "qw")) - 1) < 1e-8,
                    "Output quaternion is not unit length")
        final_tf = [row for row in fixture.transforms if row["stamp_ns"] == final_ns]
        require(bool(final_tf), "Missing final odom to base_link transform")
        require(any(all(abs(row[key] - final[key]) < 1e-9
                        for key in ("x_m", "y_m", "z_m", "qx", "qy", "qz", "qw"))
                    for row in final_tf), "Published TF differs from final filtered pose")
        save_csv(work / f"{case}-filtered.csv", fixture.rows)
        save_csv(work / f"{case}-tf.csv", fixture.transforms)
        final["position_error_m"] = math.hypot(final["x_m"] - 5.0, final["y_m"])
        completed = True
        return final
    finally:
        exit_code = stop(process, stream)
        spin_until(fixture, lambda: fixture.count_publishers("/odometry/filtered") == 0
                   and fixture.odom_pub.get_subscription_count() == 0,
                   "stopped estimator disappears from graph")
        fixture.destroy_node()
        if completed:
            require(exit_code in (0, -signal.SIGINT, 128 + signal.SIGINT),
                    f"Unexpected estimator shutdown status: {exit_code}")


def run(work):
    version = subprocess.check_output(
        ["dpkg-query", "-W", "-f=${Version}", "ros-jazzy-robot-localization"],
        text=True, timeout=5)
    require(version == PACKAGE_VERSION, f"Untested package version: {version}")
    print(f"robot_localization Debian package: {version}")
    configs = {case: yaml.safe_load((HERE / f"{case}.yaml").read_text())[
        "ekf_filter_node"]["ros__parameters"] for case in CASES}
    baseline, changed = (configs[case] for case in CASES)
    require({k: v for k, v in baseline.items() if k != "imu0_config"}
            == {k: v for k, v in changed.items() if k != "imu0_config"},
            "The two configurations differ outside the IMU selection mask")
    for config in configs.values():
        require(len(config["odom0_config"]) == len(config["imu0_config"]) == 15,
                "Selection masks must each contain 15 fields")
        require([index for index, selected in enumerate(config["odom0_config"]) if selected]
                == [6, 7, 11], "Wheel mask must select only vx, vy, and vyaw")
    require(not any(baseline["imu0_config"]), "Baseline must exclude the IMU")
    differences = [index for index in range(15)
                   if baseline["imu0_config"][index] != changed["imu0_config"][index]]
    require(differences == [11] and not baseline["imu0_config"][11]
            and changed["imu0_config"][11], "Change must enable only the IMU yaw-rate bit")
    supplied = [supplied_row(index) for index in range(STEPS + 1)]
    save_csv(work / "supplied-inputs.csv", supplied)
    results = {}
    rclpy.init()
    try:
        for case in CASES:
            results[case] = run_case(case, work)
    finally:
        rclpy.shutdown()
    base, biased = (results[case] for case in CASES)
    require(base["position_error_m"] < 0.05 and abs(base["yaw_rad"]) < 0.01,
            "Baseline differs from the supplied straight-line reference")
    require(0.15 < biased["yaw_rate_rad_s"] < 0.21,
            "Biased gyro did not influence the estimated yaw rate as expected")
    require(0.7 < biased["yaw_rad"] < 1.2 and biased["y_m"] > 1.5,
            "Biased-gyro estimate did not develop the expected positive turn")
    require(biased["position_error_m"] > 1.5,
            "Changed configuration did not worsen this supplied reference error")
    report = {
        "package_version": version,
        "fixture": {"samples_per_topic_per_case": STEPS + 1,
                    "start_ns": START_NS, "step_ns": STEP_NS,
                    "reference_final": {"x_m": 5.0, "y_m": 0.0, "yaw_rad": 0.0},
                    "odom_pose_frame": "odom", "odom_twist_frame": "base_link",
                    "imu_frame": "base_link", "supplied_inputs": supplied},
        "configurations": configs,
        "only_parameter_change": "imu0_config[11]: false -> true (vyaw)",
        "results": results,
        "checks": {"single_config_bit_changed": True, "fresh_filter_per_case": True,
                   "finite_planar_outputs": True, "unit_quaternions": True,
                   "final_tf_matches_filtered_pose": True,
                   "baseline_position_error_below_0_05_m": True,
                   "biased_position_error_above_1_5_m": True},
    }
    (work / "report.json").write_text(json.dumps(report, indent=2) + "\n")
    print("SUPPLIED: 51 odometry and 51 IMU messages per run; t=100..105 seconds")
    print("REFERENCE: x=5.000 m, y=0.000 m, yaw=0.000 rad at t=105 seconds")
    for case in CASES:
        final = results[case]
        print(f"{case}: x={final['x_m']:.3f} m, y={final['y_m']:.3f} m, "
              f"yaw={final['yaw_rad']:.3f} rad, "
              f"position error={final['position_error_m']:.3f} m")
    print("CONFIG: only imu0_config[11] changed from false to true")
    print("TF: final odom -> base_link matches filtered odometry in both runs")
    print("PASS: real ekf_node runs match the bounded fixture expectations")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, help="Retain results in a NEW directory")
    args = parser.parse_args()
    if args.output is not None:
        args.output.mkdir(parents=True, exist_ok=False)
        work = args.output.resolve()
        try:
            run(work)
        except BaseException:
            for path in sorted(work.glob("*.log")):
                print(f"\n--- {path.name} ---\n{path.read_text()}", flush=True)
            raise
        print(f"ARTIFACTS: {work}")
    else:
        with tempfile.TemporaryDirectory(prefix="localization-lesson-") as temporary:
            work = Path(temporary)
            try:
                run(work)
            except BaseException:
                for path in sorted(work.glob("*.log")):
                    print(f"\n--- {path.name} ---\n{path.read_text()}", flush=True)
                raise


if __name__ == "__main__":
    main()