#!/usr/bin/env python3
"""
wall_kf.py : Lecture 5 anchor demo. A 1D Kalman filter on the distance to a wall.

    python3 wall_kf.py --r 0.02 --q 0.01        # robot facing a wall: drive it slowly toward the wall, never in reverse
    python3 wall_kf.py --sim                    # no robot: the lecture's simulated run, live
    python3 wall_kf.py --sim --headless 25      # no window: print the RMS errors, as on the slides

State d, the distance to the wall straight ahead.
    predict, on every /odom message:  d = d - u,          P = P + q^2 |u|
    update, on every /scan:           K = P / (P + r^2),  d = d + K (z - d),  P = (1 - K) P
q is in meters per root meter of travel, r in meters.

Keys, in the plot window:
    l   LiDAR updates off and on          r / R   the filter's r divided / multiplied by 10
    q / Q   the same for q                 space   reset the filter to the current reading
"""
import argparse, math, random, statistics, sys, threading, time

class Filter:
    def __init__(self, q, r): self.q, self.r = q, r; self.d = None; self.P = 0.05**2; self.lidar = True; self.K = None
    def predict(self, u):
        if self.d is None: return
        self.d -= u; self.P += self.q**2 * abs(u) + 1e-7
    def update(self, z):
        if self.d is None: self.d = z; self.P = self.r**2; return
        if not self.lidar: return
        self.K = self.P / (self.P + self.r**2); self.d += self.K * (z - self.d); self.P *= (1 - self.K)

class Sim:
    """The lecture's run: toward the wall at 0.08 m/s for 15 s, then away; LiDAR at 5 Hz, sigma 2 cm."""
    def __init__(self, seed=11): self.rng = random.Random(seed); self.t = 0.0; self.d = 2.5
    def odom(self, dt=0.05):
        self.t += dt; v = 0.08 if (self.t % 30) < 15 else -0.08; step = v * dt; self.d -= step
        return step * (1 + self.rng.gauss(0, 0.05))
    def scan(self): return self.d + self.rng.gauss(0, 0.02)

def forward_reading(m, fwd, half):
    n = len(m.ranges); i0 = int(round((fwd - half - m.angle_min) / m.angle_increment)); i1 = int(round((fwd + half - m.angle_min) / m.angle_increment))
    good = [m.ranges[i % n] for i in range(min(i0, i1), max(i0, i1) + 1) if math.isfinite(m.ranges[i % n]) and m.range_min <= m.ranges[i % n] <= m.range_max]
    return statistics.median(good) if good else None

def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--q", type=float, default=0.01); ap.add_argument("--r", type=float, default=0.02)
    ap.add_argument("--forward-deg", type=float, default=0.0); ap.add_argument("--half-deg", type=float, default=1.0)
    ap.add_argument("--sim", action="store_true"); ap.add_argument("--headless", type=float, metavar="SECONDS")
    ap.add_argument("--window", type=float, default=30.0, help="seconds of history to plot")
    a, ros_args = ap.parse_known_args()
    f = Filter(a.q, a.r); log = []; lock = threading.Lock(); t0 = time.time(); odom_only = [None]; last_z = [None]
    def record(t, z=None, truth=None):
        if z is not None: last_z[0] = z
        with lock: log.append((t, f.d, math.sqrt(f.P), z, odom_only[0], truth))

    if a.sim and a.headless:
        sim = Sim(); f.d = sim.d + 0.04; odom_only[0] = f.d; next_scan = 0.2; T = a.headless
        while sim.t < T - 1e-9:
            u = sim.odom(); f.predict(u); odom_only[0] -= u
            z = None
            if sim.t >= next_scan - 1e-9:
                next_scan += 0.2; f.lidar = not (12.0 <= sim.t < 17.0); z = sim.scan(); f.update(z)
            record(sim.t, z, sim.d)
        e = [(row[1] - row[5]) for row in log]; eo = [(row[4] - row[5]) for row in log]
        print(f"filter RMS {100 * math.sqrt(sum(x * x for x in e) / len(e)):.2f} cm, odometry alone {100 * math.sqrt(sum(x * x for x in eo) / len(eo)):.2f} cm, "
              f"final K {f.K:.3f}, final sigma {100 * math.sqrt(f.P):.2f} cm")
        return 0

    if a.sim:
        sim = Sim(); f.d = sim.d + 0.04; odom_only[0] = f.d
        def run():
            next_scan = 0.2
            while True:
                u = sim.odom(); f.predict(u); odom_only[0] -= u; z = None
                if sim.t >= next_scan - 1e-9: next_scan += 0.2; z = sim.scan(); f.update(z)
                record(sim.t, z, sim.d); time.sleep(0.05)
        threading.Thread(target=run, daemon=True).start()
    else:
        import rclpy
        from rclpy.node import Node
        from rclpy.qos import qos_profile_sensor_data
        from nav_msgs.msg import Odometry
        from sensor_msgs.msg import LaserScan
        fwd, half = math.radians(a.forward_deg), math.radians(a.half_deg)
        class N_(Node):
            def __init__(self):
                super().__init__("wall_kf"); self.prev = None
                self.create_subscription(Odometry, "/odom", self.on_odom, qos_profile_sensor_data)
                self.create_subscription(LaserScan, "/scan", self.on_scan, qos_profile_sensor_data)
            def on_odom(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:
                    u = (p.x - self.prev[0]) * math.cos(th) + (p.y - self.prev[1]) * math.sin(th)   # forward motion = toward the wall
                    f.predict(u)
                    if odom_only[0] is not None: odom_only[0] -= u
                self.prev = (p.x, p.y); record(time.time() - t0)
            def on_scan(self, m):
                z = forward_reading(m, fwd, half)
                if z is None: return
                if odom_only[0] is None: odom_only[0] = z
                f.update(z); record(time.time() - t0, z)
                if f.K is not None and f.lidar: print(f"\r  z {z:.3f} m   estimate {f.d:.3f} m   sigma {100 * math.sqrt(f.P):.2f} cm   K {f.K:.3f}   ", end="", flush=True)
        rclpy.init(args=[sys.argv[0]] + ros_args); node = N_()
        threading.Thread(target=rclpy.spin, args=(node,), daemon=True).start()

    import matplotlib.pyplot as plt
    from matplotlib.animation import FuncAnimation
    # matplotlib's own shortcuts use l (log scale), r (home), and q (quit): free them for the demo keys
    for name in [k for k in plt.rcParams if k.startswith("keymap.")]:
        plt.rcParams[name] = [k for k in plt.rcParams[name] if k not in ("l", "r", "R", "q", "Q", " ")]
    fig, ax = plt.subplots(figsize=(11, 5)); fig.canvas.manager.set_window_title("wall_kf")
    def on_key(ev):
        if ev.key == "l": f.lidar = not f.lidar
        elif ev.key == "r": f.r /= 10
        elif ev.key == "R": f.r *= 10
        elif ev.key == "q": f.q /= 10
        elif ev.key == "Q": f.q *= 10
        elif ev.key == " " and last_z[0] is not None: f.d = last_z[0]; f.P = f.r**2; odom_only[0] = last_z[0]
    fig.canvas.mpl_connect("key_press_event", on_key)
    def draw(_):
        with lock: rows = [r for r in log if r[1] is not None]
        if not rows: return
        tnow = rows[-1][0]; rows = [r for r in rows if r[0] >= tnow - a.window]
        t = [r[0] for r in rows]; d = [r[1] for r in rows]; s = [r[2] for r in rows]
        ax.clear()
        ax.fill_between(t, [x - 2 * y for x, y in zip(d, s)], [x + 2 * y for x, y in zip(d, s)], color="#102A43", alpha=0.15, label="\u00b12\u03c3")
        zs = [(r[0], r[3]) for r in rows if r[3] is not None]
        if zs: ax.scatter(*zip(*zs), s=10, color="#0070C0", alpha=0.6, label="LiDAR")
        oo = [(r[0], r[4]) for r in rows if r[4] is not None]
        if oo: ax.plot(*zip(*oo), color="#5A6470", lw=1.5, label="odometry alone")
        ax.plot(t, d, color="#102A43", lw=2.2, label="Kalman estimate")
        if a.sim: ax.plot(t, [r[5] for r in rows], color="#2E9E4F", lw=1, ls="--", label="truth")
        ax.set_xlabel("time (s)"); ax.set_ylabel("distance to the wall (m)"); ax.legend(loc="upper right", fontsize=9, ncol=3)
        kt = f"{f.K:.3f}" if f.K is not None else "-"
        ax.set_title(f"LiDAR {'ON' if f.lidar else 'OFF'}   r {100 * f.r:.3g} cm   q {100 * f.q:.3g} cm/\u221am   K {kt}   \u03c3 {100 * math.sqrt(f.P):.2f} cm", loc="left")
    _anim = FuncAnimation(fig, draw, interval=200, cache_frame_data=False)
    plt.show()
    return 0

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