#!/usr/bin/env python3
"""
PIDtracking — 跟踪数据录制 + 可视化一体脚本

用法:
  录制:  python3 PIDtracking.py record
  制图:  python3 PIDtracking.py plot <数据目录>

输出目录: ~/yiliao_ws/datas/PIDtracking/<timestamp>/
输出文件:
  plan.csv       — 路径点 (x, y, yaw)
  cmd_vel.csv    — 控制命令 (t, vx, vz)
  odom.csv       — 里程计位姿 (t, x, y, yaw)
  goal_pose.csv  — 目标点 (x, y)
  meta.csv       — 录制参数

制图输出:
  <dir>/plots/01_trajectory.png   — 轨迹对比
  <dir>/plots/02_cmd_vel.png      — 控制命令时间序列
  <dir>/plots/03_errors.png       — 横向/航向误差
  <dir>/plots/04_curvature.png    — 曲率分析
"""

import os
import sys
import time
import math
import csv
import threading
import pathlib

# ── 录制依赖 (ros2) ──
RECORD_AVAILABLE = False
try:
    import rclpy
    from rclpy.node import Node
    from nav_msgs.msg import Path as NavPath, Odometry
    from geometry_msgs.msg import Twist, PoseStamped
    RECORD_AVAILABLE = True
except ImportError:
    pass

# ── 制图依赖 ──
PLOT_AVAILABLE = False
try:
    import numpy as np
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    PLOT_AVAILABLE = True
except ImportError:
    pass

# ============================================================================
#  录制
# ============================================================================
def mode_record(args):
    if not RECORD_AVAILABLE:
        print("ERROR: rclpy not available. Run on the robot with ROS2 sourced.")
        sys.exit(1)

    base_path = pathlib.Path.home() / "yiliao_ws" / "datas" / "PIDtracking"
    if len(args) > 2:
        base_path = pathlib.Path(args[2])
    ts = time.strftime("%Y%m%d_%H%M%S")
    out_dir = base_path / ts
    out_dir.mkdir(parents=True, exist_ok=True)

    rclpy.init()
    recorder = RecorderNode(out_dir)
    print(f"[RECORD] Output: {out_dir}")
    print(f"[RECORD] Subscribing: /plan, /cmd_vel, /goal_pose, /odom_combined (fallback /odom)")
    print(f"[RECORD] Send /goal_pose to start. Ctrl+C to stop.")
    try:
        rclpy.spin(recorder)
    except KeyboardInterrupt:
        pass
    finally:
        recorder.save_all()
        recorder.destroy_node()
        rclpy.shutdown()
    print(f"[RECORD] Saved {len(recorder.plan_data)} plans, {len(recorder.cmd_data)} cmd_vel, "
          f"{len(recorder.odom_data)} odom, {len(recorder.goal_data)} goals → {out_dir}")


class RecorderNode(Node):
    def __init__(self, out_dir):
        super().__init__("pidtracking_recorder")
        self.out_dir = out_dir
        self.lock = threading.Lock()

        self.plan_data = []   # [(t, [(x,y,yaw),...])]
        self.cmd_data = []    # [(t, vx, vz)]
        self.odom_data = []   # [(t, x, y, yaw)]
        self.goal_data = []   # [(t, x, y)]

        # 订阅
        self.plan_sub = self.create_subscription(NavPath, "/plan", self.plan_cb, 10)
        self.cmd_sub = self.create_subscription(Twist, "/cmd_vel", self.cmd_cb, 10)
        self.goal_sub = self.create_subscription(PoseStamped, "/goal_pose", self.goal_cb, 10)

        # odom: 优先 odom_combined（EKF融合），否则 /odom
        odom_topic = "/odom_combined"
        topics = self.get_topic_names_and_types()
        if "/odom_combined" not in [t[0] for t in topics]:
            odom_topic = "/odom"
            self.get_logger().info(f"/odom_combined not found, using {odom_topic}")
        self.odom_sub = self.create_subscription(Odometry, odom_topic, self.odom_cb, 10)

    def plan_cb(self, msg):
        t = time.time()
        poses = []
        for p in msg.poses:
            qx, qy, qz, qw = (p.pose.orientation.x, p.pose.orientation.y,
                               p.pose.orientation.z, p.pose.orientation.w)
            yaw = math.atan2(2.0*(qw*qz + qx*qy), 1.0 - 2.0*(qy*qy + qz*qz))
            poses.append((p.pose.position.x, p.pose.position.y, yaw))
        with self.lock:
            self.plan_data.append((t, poses))

    def cmd_cb(self, msg):
        with self.lock:
            self.cmd_data.append((time.time(), msg.linear.x, msg.angular.z))

    def odom_cb(self, msg):
        qx, qy, qz, qw = (msg.pose.pose.orientation.x, msg.pose.pose.orientation.y,
                           msg.pose.pose.orientation.z, msg.pose.pose.orientation.w)
        yaw = math.atan2(2.0*(qw*qz + qx*qy), 1.0 - 2.0*(qy*qy + qz*qz))
        with self.lock:
            self.odom_data.append((time.time(), msg.pose.pose.position.x,
                                   msg.pose.pose.position.y, yaw))

    def goal_cb(self, msg):
        with self.lock:
            self.goal_data.append((time.time(), msg.pose.position.x, msg.pose.position.y))
        self.get_logger().info(f"Goal: ({msg.pose.position.x:.2f}, {msg.pose.position.y:.2f})")

    def save_all(self):
        with self.lock:
            # plan
            with open(self.out_dir / "plan.csv", "w", newline="") as f:
                w = csv.writer(f)
                w.writerow(["t", "x", "y", "yaw"])
                for t, poses in self.plan_data:
                    for x, y, yaw in poses:
                        w.writerow([t, x, y, yaw])

            # cmd_vel
            with open(self.out_dir / "cmd_vel.csv", "w", newline="") as f:
                w = csv.writer(f)
                w.writerow(["t", "vx", "vz"])
                for t, vx, vz in self.cmd_data:
                    w.writerow([t, vx, vz])

            # odom
            with open(self.out_dir / "odom.csv", "w", newline="") as f:
                w = csv.writer(f)
                w.writerow(["t", "x", "y", "yaw"])
                for t, x, y, yaw in self.odom_data:
                    w.writerow([t, x, y, yaw])

            # goal
            with open(self.out_dir / "goal_pose.csv", "w", newline="") as f:
                w = csv.writer(f)
                w.writerow(["t", "x", "y"])
                for t, x, y in self.goal_data:
                    w.writerow([t, x, y])

            # meta
            with open(self.out_dir / "meta.csv", "w", newline="") as f:
                w = csv.writer(f)
                w.writerow(["key", "value"])
                w.writerow(["plans_received", len(self.plan_data)])
                w.writerow(["cmd_vel_msgs", len(self.cmd_data)])
                w.writerow(["odom_msgs", len(self.odom_data)])
                w.writerow(["goals_received", len(self.goal_data)])


# ============================================================================
#  制图
# ============================================================================
def mode_plot(args):
    if not PLOT_AVAILABLE:
        print("ERROR: numpy/matplotlib not available. pip install numpy matplotlib")
        sys.exit(1)

    if len(args) < 3:
        print("Usage: python3 PIDtracking.py plot <data_dir>")
        sys.exit(1)

    data_dir = pathlib.Path(args[2])
    if not data_dir.exists():
        print(f"ERROR: Directory not found: {data_dir}")
        sys.exit(1)

    out_dir = data_dir / "plots"
    out_dir.mkdir(exist_ok=True)

    # 读取数据
    plan = _read_plan(data_dir / "plan.csv")
    cmd_vel = _read_csv(data_dir / "cmd_vel.csv", ["t", "vx", "vz"])
    odom = _read_csv(data_dir / "odom.csv", ["t", "x", "y", "yaw"])
    goals = _read_csv(data_dir / "goal_pose.csv", ["t", "x", "y"])

    print(f"Data: plan={len(plan)} paths, cmd_vel={len(cmd_vel)} msgs, odom={len(odom)} msgs")

    if len(odom) == 0:
        print("ERROR: No odometry data")
        sys.exit(1)

    t0 = odom[0]["t"]

    # ── 图1: 轨迹 ──
    _plot_trajectory(plan, odom, goals, out_dir)

    # ── 图2: cmd_vel ──
    if cmd_vel:
        _plot_cmd_vel(cmd_vel, t0, out_dir)

    # ── 图3: 误差 ──
    if plan and odom:
        _plot_errors(plan, odom, t0, out_dir)

    # ── 图4: 曲率 ──
    if plan and cmd_vel:
        _plot_curvature(plan, out_dir)

    print(f"\nPlots saved to: {out_dir}/")


def _read_csv(path, columns):
    if not path.exists():
        return []
    data = []
    with open(path, "r") as f:
        reader = csv.DictReader(f)
        for row in reader:
            item = {}
            for col in columns:
                item[col] = float(row[col])
            data.append(item)
    return data


def _read_plan(path):
    """读取 plan.csv: t,x,y,yaw，按 t 分组"""
    if not path.exists():
        return []
    plans = []
    current_t = None
    current_poses = []
    with open(path, "r") as f:
        reader = csv.DictReader(f)
        for row in reader:
            t = float(row["t"])
            x, y, yaw = float(row["x"]), float(row["y"]), float(row["yaw"])
            if current_t is None:
                current_t = t
            if abs(t - current_t) > 0.01:  # 新的一组
                plans.append({"t": current_t, "poses": current_poses})
                current_t = t
                current_poses = []
            current_poses.append({"x": x, "y": y, "yaw": yaw})
        if current_poses:
            plans.append({"t": current_t, "poses": current_poses})
    return plans


def _plot_trajectory(plans, odom, goals, out_dir):
    fig, ax = plt.subplots(figsize=(10, 10))

    # 最新 plan
    if plans:
        last = plans[-1]["poses"]
        if last:
            px = [p["x"] for p in last]
            py = [p["y"] for p in last]
            ax.plot(px, py, "b--", lw=1.5, alpha=0.6, label="Plan")
            ax.scatter(px[0], py[0], c="blue", s=80, marker="o", zorder=5, label="Plan Start")
            ax.scatter(px[-1], py[-1], c="blue", s=80, marker="x", zorder=5, label="Plan Goal")

    # odom
    ox = np.array([d["x"] for d in odom])
    oy = np.array([d["y"] for d in odom])
    oyaw = np.array([d["yaw"] for d in odom])
    ax.plot(ox, oy, "r-", lw=1.2, alpha=0.85, label="Actual")
    ax.scatter(ox[0], oy[0], c="red", s=60, marker="o", zorder=5, label="Start")
    ax.scatter(ox[-1], oy[-1], c="red", s=60, marker="x", zorder=5, label="End")

    # 航向箭头
    step = max(1, len(ox) // 25)
    for i in range(0, len(ox), step):
        dx = 0.06 * math.cos(oyaw[i])
        dy = 0.06 * math.sin(oyaw[i])
        ax.arrow(ox[i], oy[i], dx, dy, head_width=0.04, head_length=0.04,
                 fc="orange", ec="orange", alpha=0.5, zorder=6)

    # goal
    for g in goals:
        ax.scatter(g["x"], g["y"], c="green", s=120, marker="*", zorder=7,
                   edgecolors="darkgreen", linewidths=0.5)
    if goals:
        ax.scatter([], [], c="green", s=80, marker="*", label="Goal")

    ax.set_xlabel("X (m)")
    ax.set_ylabel("Y (m)")
    ax.set_title("Trajectory: Planned vs Actual")
    ax.legend()
    ax.axis("equal")
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    fig.savefig(out_dir / "01_trajectory.png", dpi=150)
    plt.close(fig)
    print("  Saved: 01_trajectory.png")


def _plot_cmd_vel(cmd_vel, t0, out_dir):
    t = np.array([d["t"] for d in cmd_vel]) - t0
    vx = np.array([d["vx"] for d in cmd_vel])
    vz = np.array([d["vz"] for d in cmd_vel])

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 8), sharex=True)

    ax1.plot(t, vx, "b-", lw=0.7)
    ax1.set_ylabel("Linear X (m/s)")
    ax1.set_title("cmd_vel — Linear Velocity")
    ax1.axhline(y=0, c="gray", ls=":", lw=0.5)
    ax1.grid(True, alpha=0.3)

    ax2.plot(t, vz, "r-", lw=0.7)
    ax2.set_xlabel("Time (s)")
    ax2.set_ylabel("Angular Z (rad/s)")
    ax2.set_title("cmd_vel — Angular Velocity")
    ax2.axhline(y=0, c="gray", ls=":", lw=0.5)
    ax2.grid(True, alpha=0.3)

    fig.tight_layout()
    fig.savefig(out_dir / "02_cmd_vel.png", dpi=150)
    plt.close(fig)
    print("  Saved: 02_cmd_vel.png")


def _plot_errors(plans, odom, t0, out_dir):
    """横向误差 + 航向误差"""
    ox = np.array([d["x"] for d in odom])
    oy = np.array([d["y"] for d in odom])
    oyaw = np.array([d["yaw"] for d in odom])
    ot = np.array([d["t"] for d in odom]) - t0

    # 取最新的 plan
    poses = plans[-1]["poses"]
    if len(poses) < 2:
        print("  WARNING: plan has < 2 waypoints, skipping error analysis")
        return

    pts = np.array([(p["x"], p["y"]) for p in poses])

    cte = np.zeros(len(ox))
    heading_err = np.zeros(len(ox))

    for i in range(len(ox)):
        # 横向误差
        diffs = pts[1:] - pts[:-1]
        seg_sq = np.sum(diffs**2, axis=1)
        pos = np.array([ox[i], oy[i]])
        t_proj = np.clip(np.sum((pos - pts[:-1]) * diffs, axis=1) / np.maximum(seg_sq, 1e-9), 0, 1)
        proj = pts[:-1] + t_proj[:, np.newaxis] * diffs
        dists = np.linalg.norm(proj - pos, axis=1)
        cte[i] = np.min(dists)

        # 航向误差
        idx = np.argmin(dists)
        target_idx = min(idx + 3, len(poses) - 1)
        target_yaw = poses[target_idx]["yaw"]
        err = target_yaw - oyaw[i]
        while err > math.pi:  err -= 2*math.pi
        while err < -math.pi: err += 2*math.pi
        heading_err[i] = err

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 8), sharex=True)

    ax1.plot(ot, cte, "purple", lw=0.7)
    ax1.set_ylabel("Cross-Track Error (m)")
    ax1.set_title("Cross-Track Error")
    ax1.axhline(y=0, c="gray", ls=":", lw=0.5)
    ax1.grid(True, alpha=0.3)

    ax2.plot(ot, np.degrees(heading_err), "orange", lw=0.7)
    ax2.set_xlabel("Time (s)")
    ax2.set_ylabel("Heading Error (deg)")
    ax2.set_title("Heading Error")
    ax2.axhline(y=0, c="gray", ls=":", lw=0.5)
    ax2.grid(True, alpha=0.3)

    fig.tight_layout()
    fig.savefig(out_dir / "03_errors.png", dpi=150)
    plt.close(fig)
    print("  Saved: 03_errors.png")


def _plot_curvature(plans, out_dir):
    poses = plans[-1]["poses"]
    n = len(poses)
    if n < 3:
        return

    # 累积路径距离 (沿路径从 pose[0] 起算)
    cum_dist = [0.0]
    for i in range(1, n):
        dx = poses[i]["x"] - poses[i-1]["x"]
        dy = poses[i]["y"] - poses[i-1]["y"]
        cum_dist.append(cum_dist[-1] + math.hypot(dx, dy))

    # 曲率 (三点法，每个曲率值对应中间点的路径位置)
    curvatures = []
    curv_dist = []
    for i in range(1, n - 1):
        x0, y0 = poses[i-1]["x"], poses[i-1]["y"]
        x1, y1 = poses[i]["x"], poses[i]["y"]
        x2, y2 = poses[i+1]["x"], poses[i+1]["y"]
        a = math.hypot(x1-x0, y1-y0)
        b = math.hypot(x2-x1, y2-y1)
        c = math.hypot(x2-x0, y2-y0)
        if a*b*c < 1e-12:
            curvatures.append(0.0)
        else:
            s = (a+b+c)/2.0
            area = math.sqrt(max(0, s*(s-a)*(s-b)*(s-c)))
            curvatures.append(4.0*area/(a*b*c))
        curv_dist.append(cum_dist[i])  # 曲率对应 pose[i] 在路径上的位置

    dist_mid = np.array(curv_dist)
    curv_arr = np.array(curvatures)

    fig, ax = plt.subplots(figsize=(12, 5))
    ax.plot(dist_mid, curv_arr, "g-", lw=1.0)
    ax.set_xlabel("Distance along path (m)")
    ax.set_ylabel("Curvature (1/m)")
    ax.set_title("Path Curvature vs Distance")
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    fig.savefig(out_dir / "04_curvature.png", dpi=150)
    plt.close(fig)
    print("  Saved: 04_curvature.png")


# ============================================================================
#  main
# ============================================================================
if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)

    cmd = sys.argv[1].lower()

    if cmd == "record":
        mode_record(sys.argv)
    elif cmd == "plot":
        mode_plot(sys.argv)
    else:
        print(f"Unknown command: {cmd}")
        print("Use: record  or  plot <dir>")
        sys.exit(1)