#!/usr/bin/env python3
"""
cov_odom.py : Lecture 5 Section 2 demo. Odometry's uncertainty, propagated live.

    python3 cov_odom.py --k 1e-5             # then in RViz: add a Marker display on /cov_odom/ellipse
    python3 cov_odom.py --selftest           # check the propagation against the lecture's numbers

Reads /odom pose increments, turns each into left and right wheel distances,
and propagates the 3 x 3 covariance with the Lecture 5 formula:

    P' = Fx P Fx^T + Fu Q Fu^T,   Q = diag(k |d_r|, k |d_l|)

It draws the 2-sigma position ellipse at the robot, in the odom frame, and
prints the sigmas once a second. Press r then Enter to reset to a known pose.

/odom is used instead of /wheel_ticks because on Jazzy /wheel_ticks is only
exposed under /_do_not_use on the robot. The propagation is the same.
"""
import argparse, math, select, sys, time
import numpy as np

W = 0.235

def step(P, theta, dr, dl, k, w=W):
    d = 0.5 * (dr + dl); dth = (dr - dl) / w; ph = theta + dth / 2; c, s = math.cos(ph), math.sin(ph)
    Fx = np.array([[1, 0, -d * s], [0, 1, d * c], [0, 0, 1]])
    Fu = np.array([[0.5 * c - d / (2 * w) * s, 0.5 * c + d / (2 * w) * s],
                   [0.5 * s + d / (2 * w) * c, 0.5 * s - d / (2 * w) * c],
                   [1 / w, -1 / w]])
    return Fx @ P @ Fx.T + Fu @ np.diag([k * abs(dr), k * abs(dl)]) @ Fu.T

def ellipse_points(P2, k_sigma=2.0, n=48):
    vals, vecs = np.linalg.eigh(P2); vals = np.clip(vals, 0, None)
    t = np.linspace(0, 2 * math.pi, n + 1)
    circle = np.vstack([np.cos(t), np.sin(t)])
    return (vecs @ np.diag(k_sigma * np.sqrt(vals)) @ circle).T

def selftest():
    P = np.zeros((3, 3)); th = 0.0; out = {}
    for i in range(1, 401):
        P = step(P, th, 0.01, 0.01, 1e-5)
        if i in (100, 200, 400): out[i / 100] = (100 * math.sqrt(P[0, 0]), 100 * math.sqrt(P[1, 1]), math.degrees(math.sqrt(P[2, 2])))
    want = {1.0: (0.22, 1.10, 1.09), 2.0: (0.32, 3.11, 1.54), 4.0: (0.45, 8.79, 2.18)}
    ok = all(abs(out[s][j] - want[s][j]) < 0.01 for s in want for j in range(3))
    for s in want:
        print(f"  {s:.0f} m: sigma x {out[s][0]:.2f} cm, y {out[s][1]:.2f} cm, theta {out[s][2]:.2f} deg   (slide: {want[s]})")
    pts = ellipse_points(np.diag([0.04**2, 0.01**2]))
    ok &= abs(pts[:, 0].max() - 0.08) < 1e-9 and abs(pts[:, 1].max() - 0.02) < 1e-3
    print("ALL CHECKS PASSED" if ok else "SOME CHECKS FAILED"); return 0 if ok else 1

def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--k", type=float, default=1e-5, help="wheel variance per meter rolled, m")
    ap.add_argument("--topic", default="/odom"); ap.add_argument("--selftest", action="store_true")
    a, ros_args = ap.parse_known_args()
    if a.selftest: return selftest()
    import rclpy
    from rclpy.node import Node
    from rclpy.qos import qos_profile_sensor_data
    from nav_msgs.msg import Odometry
    from visualization_msgs.msg import Marker
    from geometry_msgs.msg import Point
    class C(Node):
        def __init__(self):
            super().__init__("cov_odom"); self.P = np.zeros((3, 3)); self.prev = None; self.t = 0.0
            self.pub = self.create_publisher(Marker, "/cov_odom/ellipse", 10)
            self.create_subscription(Odometry, a.topic, self.cb, qos_profile_sensor_data)
            self.create_timer(0.2, self.keys)
            self.get_logger().info("RViz: fixed frame odom, add Marker on /cov_odom/ellipse. Type r + Enter to reset.")
        def keys(self):
            if select.select([sys.stdin], [], [], 0)[0] and sys.stdin.readline().strip() == "r":
                self.P = np.zeros((3, 3)); self.get_logger().info("reset: the pose is known exactly again")
        def cb(self, m):
            p = m.pose.pose.position; q = m.pose.pose.orientation
            th = math.atan2(2 * (q.w * q.z + q.x * q.y), 1 - 2 * (q.y * q.y + q.z * q.z))
            if self.prev is not None:
                x0, y0, t0 = self.prev
                dth = math.atan2(math.sin(th - t0), math.cos(th - t0))
                d = (p.x - x0) * math.cos(t0 + dth / 2) + (p.y - y0) * math.sin(t0 + dth / 2)
                self.P = step(self.P, t0, d + dth * W / 2, d - dth * W / 2, a.k)
            self.prev = (p.x, p.y, th)
            mk = Marker(); mk.header = m.header; mk.ns = "cov_odom"; mk.id = 0; mk.type = Marker.LINE_STRIP; mk.action = Marker.ADD
            mk.scale.x = 0.01; mk.color.r, mk.color.g, mk.color.b, mk.color.a = 0.0, 0.44, 0.75, 1.0; mk.pose.orientation.w = 1.0
            for ex, ey in ellipse_points(self.P[:2, :2]):
                mk.points.append(Point(x=p.x + float(ex), y=p.y + float(ey), z=0.02))
            self.pub.publish(mk)
            if time.time() - self.t > 1.0:
                self.t = time.time()
                self.get_logger().info(f"sigma x {100 * math.sqrt(self.P[0, 0]):.2f} cm, y {100 * math.sqrt(self.P[1, 1]):.2f} cm, "
                                       f"theta {math.degrees(math.sqrt(self.P[2, 2])):.2f} deg (in the odom frame)")
    rclpy.init(args=[sys.argv[0]] + ros_args); node = C()
    try: rclpy.spin(node)
    except KeyboardInterrupt: pass
    finally:
        node.destroy_node()
        if rclpy.ok(): rclpy.shutdown()
    return 0

if __name__ == "__main__":
    sys.exit(main())