#!/usr/bin/env python3
"""
Measure actual turning radius at different speeds for Ackermann robot.

Usage (on RDKx5):
    python3 measure_turning_radius.py [speed1 speed2 ...]

    Default speeds: 0.1 0.2 0.3 0.4 0.5

Saves data to: ~/turning_radius_data/YYYYMMDD_HHMMSS/
    raw_v{速度}.csv      — 每条 odom 原始数据
    summary.csv          — 汇总结果
    trajectory_v{速度}.png — 轨迹图 (可选)

Safety:
    - Large open space required (at least 2x the max turning radius)
    - Ctrl+C to abort at any time
    - Robot stops automatically after each test
"""

import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Twist
from nav_msgs.msg import Odometry
import math
import sys
import time
import os
import csv
from datetime import datetime
import numpy as np


class TurningRadiusMeasurer(Node):
    def __init__(self, speeds, data_dir, max_ang=5.0):
        super().__init__("turning_radius_measure")
        self.pub = self.create_publisher(Twist, "/cmd_vel", 10)
        self.odom_sub = self.create_subscription(Odometry, "/odom", self.odom_cb, 10)

        self.speeds = speeds
        self.max_ang = max_ang
        self.data_dir = data_dir
        self.results = []

        self._poses = []
        self._v_samples = []
        self._w_samples = []
        self._timestamps = []
        self._collecting = False
        self._start_yaw = None

    def odom_cb(self, msg):
        if not self._collecting:
            return
        pos = msg.pose.pose.position
        quat = msg.pose.pose.orientation
        yaw = self._quat_to_yaw(quat)
        v = msg.twist.twist.linear.x
        w = msg.twist.twist.angular.z
        t = msg.header.stamp.sec + msg.header.stamp.nanosec * 1e-9

        self._poses.append((pos.x, pos.y, yaw))
        self._v_samples.append(v)
        self._w_samples.append(w)
        self._timestamps.append(t)

        if self._start_yaw is None:
            self._start_yaw = yaw

    def _quat_to_yaw(self, q):
        siny = 2.0 * (q.w * q.z + q.x * q.y)
        cosy = 1.0 - 2.0 * (q.y * q.y + q.z * q.z)
        return math.atan2(siny, cosy)

    def _unwrap_delta(self, poses):
        deltas = []
        for i in range(1, len(poses)):
            d = poses[i][2] - poses[i-1][2]
            d = math.atan2(math.sin(d), math.cos(d))
            deltas.append(d)
        return deltas

    def _total_angle(self):
        if len(self._poses) < 2:
            return 0.0
        deltas = self._unwrap_delta(self._poses)
        return sum(abs(d) for d in deltas)

    def run_test(self, lin_speed):
        self._poses = []
        self._v_samples = []
        self._w_samples = []
        self._timestamps = []
        self._collecting = False
        self._start_yaw = None

        ang_speed = self.max_ang

        self.get_logger().info(
            "Testing v=%.2f m/s, ω=%.2f rad/s" % (lin_speed, ang_speed))
        self.get_logger().info("Starting in 2 seconds...")
        time.sleep(2)

        twist = Twist()
        twist.linear.x = lin_speed
        twist.angular.z = ang_speed
        self.pub.publish(twist)

        self._collecting = True
        self.get_logger().info("Collecting...")

        while rclpy.ok():
            rclpy.spin_once(self, timeout_sec=0.05)
            if self._total_angle() > 6.0:
                break

        self._collecting = False
        twist.linear.x = 0.0
        twist.angular.z = 0.0
        self.pub.publish(twist)

        self._compute_result(lin_speed, ang_speed)

    def _compute_result(self, lin_speed, ang_speed):
        if len(self._v_samples) < 10:
            self.get_logger().warn("Not enough data points!")
            return

        # Save raw data
        self._save_raw_data(lin_speed)

        # Compute radius
        v_trim = self._v_samples[10:]
        w_trim = self._w_samples[10:]

        v_mean = sum(abs(v) for v in v_trim) / len(v_trim)
        w_mean = sum(abs(w) for w in w_trim) / len(w_trim)

        if w_mean < 0.01:
            self.get_logger().warn("Angular velocity too low, can't compute radius")
            return

        R_vw = v_mean / w_mean
        R_fit = self._fit_circle_radius()

        # Print
        self.get_logger().info("=" * 50)
        self.get_logger().info("v_cmd=%.2f  ω_cmd=%.2f" % (lin_speed, ang_speed))
        self.get_logger().info("v_avg=%.3f  ω_avg=%.3f" % (v_mean, w_mean))
        self.get_logger().info("R (v/ω):  %.3f m" % R_vw)
        if R_fit is not None:
            self.get_logger().info("R (fit):  %.3f m" % R_fit)
        self.get_logger().info("Data saved: %s" % self.data_dir)
        self.get_logger().info("=" * 50)

        self.results.append({
            "v_cmd": lin_speed, "w_cmd": ang_speed,
            "v_avg": v_mean, "w_avg": w_mean,
            "R_vw": R_vw, "R_fit": R_fit if R_fit else -1,
            "n_samples": len(v_trim)
        })

        # Try plot
        self._try_plot(lin_speed)

    def _save_raw_data(self, lin_speed):
        """Save raw odometry to CSV"""
        fname = os.path.join(self.data_dir, "raw_v%.2f.csv" % lin_speed)
        with open(fname, "w", newline="") as f:
            w = csv.writer(f)
            w.writerow(["t", "x", "y", "yaw", "v_linear", "v_angular"])
            for i in range(len(self._poses)):
                x, y, yaw = self._poses[i]
                v = self._v_samples[i]
                ang = self._w_samples[i]
                t = self._timestamps[i] - self._timestamps[0] if self._timestamps else i
                w.writerow(["%.4f" % t, "%.6f" % x, "%.6f" % y,
                           "%.6f" % yaw, "%.6f" % v, "%.6f" % ang])
        self.get_logger().info("  raw data: %s (%d points)" % (fname, len(self._poses)))

    def _try_plot(self, lin_speed):
        """Try to save a trajectory plot (skips if matplotlib not available)"""
        try:
            import matplotlib
            matplotlib.use("Agg")
            import matplotlib.pyplot as plt

            xs = [p[0] for p in self._poses]
            ys = [p[1] for p in self._poses]

            fig, ax = plt.subplots(figsize=(6, 6))
            ax.plot(xs, ys, "b-", linewidth=0.8, label="trajectory")
            ax.plot(xs[0], ys[0], "go", label="start")
            ax.plot(xs[-1], ys[-1], "ro", label="end")
            ax.set_aspect("equal")
            ax.set_xlabel("x (m)")
            ax.set_ylabel("y (m)")
            ax.set_title("v=%.2f m/s  (%.1f° turn)" %
                         (lin_speed, math.degrees(self._total_angle())))
            ax.legend()
            ax.grid(True, alpha=0.3)

            fname = os.path.join(self.data_dir, "trajectory_v%.2f.png" % lin_speed)
            fig.savefig(fname, dpi=120)
            plt.close(fig)
            self.get_logger().info("  plot saved: %s" % fname)
        except ImportError:
            pass

    def save_summary(self):
        """Save summary CSV"""
        fname = os.path.join(self.data_dir, "summary.csv")
        with open(fname, "w", newline="") as f:
            w = csv.writer(f)
            w.writerow(["v_cmd", "w_cmd", "v_avg", "w_avg",
                        "R_vw_m", "R_fit_m", "n_samples"])
            for r in self.results:
                w.writerow([r["v_cmd"], r["w_cmd"], r["v_avg"], r["w_avg"],
                           r["R_vw"], r["R_fit"], r["n_samples"]])
        print("\nSummary saved: %s" % fname)
        print("Raw data dir:  %s" % self.data_dir)

    def _fit_circle_radius(self):
        if len(self._poses) < 10:
            return None

        xs = np.array([p[0] for p in self._poses])
        ys = np.array([p[1] for p in self._poses])

        A = np.column_stack([xs, ys, np.ones_like(xs)])
        b = -(xs**2 + ys**2)
        try:
            sol, _, _, _ = np.linalg.lstsq(A, b, rcond=None)
            a, b_coef, c = sol
            r_sq = (a**2 + b_coef**2) / 4.0 - c
            if r_sq > 0:
                return float(np.sqrt(r_sq))
        except np.linalg.LinAlgError:
            pass
        return None

    def print_summary(self):
        print("\n" + "=" * 65)
        print("  TURNING RADIUS SUMMARY")
        print("=" * 65)
        print("  v_cmd   ω_cmd    v_avg    ω_avg    R(v/ω)     R(fit)")
        print("-" * 65)
        for r in self.results:
            rf = "%.3f" % r["R_fit"] if r["R_fit"] > 0 else "N/A"
            print("  %5.2f   %5.2f   %6.3f   %6.3f   %7.3f   %s" %
                  (r["v_cmd"], r["w_cmd"], r["v_avg"], r["w_avg"],
                   r["R_vw"], rf))
        print("=" * 65)


def main():
    rclpy.init()

    speeds = [0.1, 0.15, 0.2, 0.25, 0.3, 0.35]
    if len(sys.argv) > 1:
        try:
            speeds = [float(a) for a in sys.argv[1:]]
        except ValueError:
            print("Usage: python3 measure_turning_radius.py [v1 v2 v3 ...]")
            return

    # Create data dir
    data_dir = os.path.join(
        os.path.expanduser("~/yiliao_ws"),
        "turning_radius_data",
        datetime.now().strftime("%Y%m%d_%H%M%S"))
    os.makedirs(data_dir, exist_ok=True)

    print("=" * 60)
    print("  Turning Radius Measurement")
    print("=" * 60)
    print("  Speeds:     %s" % speeds)
    print("  Max steer:  5.0 rad/s")
    print("  Data dir:   %s" % data_dir)
    print("=" * 60)
    print("  WARNING: Large open space required (2m+)!")
    input("  Press ENTER to start, Ctrl+C to abort...")

    node = TurningRadiusMeasurer(speeds, data_dir, max_ang=5.0)

    try:
        for v in speeds:
            node.run_test(v)
            time.sleep(0.5)

        node.print_summary()
        node.save_summary()

    except KeyboardInterrupt:
        twist = Twist()
        node.pub.publish(twist)
        print("\nAborted.")
        node.print_summary()
        node.save_summary()

    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()
