#!/usr/bin/env python3
"""Record and summarize real-robot planner and tracker performance."""

import csv
import json
import math
from pathlib import Path
import statistics
import time

import rclpy
from ackermann_msgs.msg import AckermannDriveStamped
from geometry_msgs.msg import PoseStamped, Twist
from nav_msgs.msg import OccupancyGrid, Odometry, Path as NavPath
from rclpy.duration import Duration
from rclpy.node import Node
from rclpy.qos import DurabilityPolicy, QoSProfile, ReliabilityPolicy
from rclpy.time import Time
from std_msgs.msg import Bool, String
from tf2_ros import Buffer, TransformException, TransformListener


def yaw_from_quaternion(q):
    return math.atan2(
        2.0 * (q.w * q.z + q.x * q.y),
        1.0 - 2.0 * (q.y * q.y + q.z * q.z),
    )


def normalize_angle(angle):
    return math.atan2(math.sin(angle), math.cos(angle))


def percentile(values, fraction):
    if not values:
        return None
    ordered = sorted(values)
    index = min(len(ordered) - 1, max(0, math.ceil(fraction * len(ordered)) - 1))
    return ordered[index]


def point_segment_distance(px, py, ax, ay, bx, by):
    dx = bx - ax
    dy = by - ay
    length_sq = dx * dx + dy * dy
    if length_sq <= 1e-12:
        return math.hypot(px - ax, py - ay)
    ratio = ((px - ax) * dx + (py - ay) * dy) / length_sq
    ratio = max(0.0, min(1.0, ratio))
    cx = ax + ratio * dx
    cy = ay + ratio * dy
    return math.hypot(px - cx, py - cy)


class NavigationTestRecorder(Node):
    def __init__(self):
        super().__init__("navigation_test_recorder")
        self.declare_parameter("test_name", "navigation_test")
        self.declare_parameter("output_dir", "~/.ros/navigation_tests")
        self.declare_parameter("map_frame", "map")
        self.declare_parameter("map_topic", "/map")
        self.declare_parameter("odom_topic", "/odom_combined")
        self.declare_parameter("goal_pose_topic", "/goal_pose")
        self.declare_parameter("plan_topic", "/plan")
        self.declare_parameter("cmd_vel_topic", "/cmd_vel")
        self.declare_parameter("drive_command_topic", "/ackermann_cmd")
        self.declare_parameter("cancel_topic", "/navigation_cancel")
        self.declare_parameter("planner_status_topic", "/hybrid_astar_status")
        self.declare_parameter("tracker_status_topic", "/pure_pursuit_status")
        self.declare_parameter("sample_rate", 20.0)
        self.declare_parameter("tf_timeout", 0.10)
        self.declare_parameter("occupied_threshold", 50)

        self.test_name = str(self.get_parameter("test_name").value)
        output_dir = Path(str(self.get_parameter("output_dir").value)).expanduser()
        output_dir.mkdir(parents=True, exist_ok=True)
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        self.output_base = output_dir / f"{self.test_name}_{timestamp}"
        self.csv_path = self.output_base.with_suffix(".csv")
        self.summary_path = self.output_base.with_suffix(".json")
        self.map_frame = str(self.get_parameter("map_frame").value)
        self.map_topic = str(self.get_parameter("map_topic").value)
        self.odom_topic = str(self.get_parameter("odom_topic").value)
        self.goal_pose_topic = str(self.get_parameter("goal_pose_topic").value)
        self.plan_topic = str(self.get_parameter("plan_topic").value)
        self.cmd_vel_topic = str(self.get_parameter("cmd_vel_topic").value)
        self.drive_command_topic = str(
            self.get_parameter("drive_command_topic").value
        )
        self.cancel_topic = str(self.get_parameter("cancel_topic").value)
        self.planner_status_topic = str(
            self.get_parameter("planner_status_topic").value
        )
        self.tracker_status_topic = str(
            self.get_parameter("tracker_status_topic").value
        )
        self.tf_timeout = float(self.get_parameter("tf_timeout").value)
        self.occupied_threshold = int(self.get_parameter("occupied_threshold").value)

        self.started_at = self.get_clock().now()
        self.latest_odom = None
        self.latest_cmd = Twist()
        self.latest_drive = AckermannDriveStamped()
        self.latest_path = []
        self.latest_goal = None
        self.map_msg = None
        self.planner_status = ""
        self.tracker_status = ""
        self.pending_goal_time = None
        self.pending_goal_xy = None
        self.plan_latencies = []
        self.cancel_stop_latencies = []
        self.cancel_pending_time = None
        self.plan_count = 0
        self.plan_ok_count = 0
        self.plan_fail_count = 0
        self.plan_collision_cells = []
        self.plan_min_turning_radii = []
        self.plan_lengths = []
        self.reverse_fractions = []
        self.cross_track_errors = []
        self.goal_errors = []
        self.samples = 0
        self.transform_failures = 0
        self.max_abs_cmd_linear = 0.0
        self.max_abs_cmd_angular = 0.0
        self.max_abs_drive_speed = 0.0
        self.max_abs_steering_angle = 0.0
        self.max_abs_measured_linear = 0.0
        self.max_abs_measured_angular = 0.0
        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, self)

        self.csv_file = self.csv_path.open("w", newline="", encoding="utf-8")
        self.csv_writer = csv.writer(self.csv_file)
        self.csv_writer.writerow([
            "elapsed_sec", "map_x", "map_y", "map_yaw",
            "measured_linear", "measured_angular",
            "cmd_linear", "cmd_angular", "drive_speed", "steering_angle",
            "cross_track_error",
            "goal_error", "plan_points", "planner_status", "tracker_status",
        ])

        map_qos = QoSProfile(
            depth=1,
            reliability=ReliabilityPolicy.RELIABLE,
            durability=DurabilityPolicy.TRANSIENT_LOCAL,
        )
        self.input_subscriptions = [
            self.create_subscription(
                OccupancyGrid, self.map_topic, self._map_cb, map_qos
            ),
            self.create_subscription(Odometry, self.odom_topic, self._odom_cb, 20),
            self.create_subscription(
                PoseStamped, self.goal_pose_topic, self._goal_cb, map_qos
            ),
            self.create_subscription(NavPath, self.plan_topic, self._plan_cb, 10),
            self.create_subscription(Twist, self.cmd_vel_topic, self._cmd_cb, 20),
            self.create_subscription(
                AckermannDriveStamped, self.drive_command_topic, self._drive_cb, 20
            ),
            self.create_subscription(Bool, self.cancel_topic, self._cancel_cb, 10),
            self.create_subscription(
                String, self.planner_status_topic, self._planner_status_cb, 20
            ),
            self.create_subscription(
                String, self.tracker_status_topic, self._tracker_status_cb, 20
            ),
        ]
        sample_rate = max(1.0, float(self.get_parameter("sample_rate").value))
        self.timer = self.create_timer(1.0 / sample_rate, self._sample)
        self.get_logger().info(f"recording navigation test to {self.csv_path}")

    def _map_cb(self, msg):
        self.map_msg = msg

    def _odom_cb(self, msg):
        self.latest_odom = msg

    def _cmd_cb(self, msg):
        self.latest_cmd = msg
        self.max_abs_cmd_linear = max(self.max_abs_cmd_linear, abs(msg.linear.x))
        self.max_abs_cmd_angular = max(self.max_abs_cmd_angular, abs(msg.angular.z))
        if (
            self.cancel_pending_time is not None
            and abs(msg.linear.x) <= 0.01
            and abs(msg.angular.z) <= 0.01
        ):
            latency = (self.get_clock().now() - self.cancel_pending_time).nanoseconds * 1e-9
            self.cancel_stop_latencies.append(max(0.0, latency))
            self.cancel_pending_time = None

    def _cancel_cb(self, msg):
        if not msg.data:
            return
        if abs(self.latest_cmd.linear.x) <= 0.01 and abs(self.latest_cmd.angular.z) <= 0.01:
            self.cancel_stop_latencies.append(0.0)
        else:
            self.cancel_pending_time = self.get_clock().now()

    def _drive_cb(self, msg):
        self.latest_drive = msg
        self.max_abs_drive_speed = max(
            self.max_abs_drive_speed, abs(msg.drive.speed)
        )
        self.max_abs_steering_angle = max(
            self.max_abs_steering_angle, abs(msg.drive.steering_angle)
        )

    def _goal_cb(self, msg):
        goal_xy = (msg.pose.position.x, msg.pose.position.y)
        if self.pending_goal_xy is None or math.hypot(
            goal_xy[0] - self.pending_goal_xy[0],
            goal_xy[1] - self.pending_goal_xy[1],
        ) > 1e-4:
            self.pending_goal_time = self.get_clock().now()
            self.pending_goal_xy = goal_xy
        self.latest_goal = goal_xy

    def _planner_status_cb(self, msg):
        self.planner_status = msg.data
        if msg.data.startswith("plan_ok"):
            self.plan_ok_count += 1
        elif msg.data.startswith("plan_failed") or msg.data.startswith("plan_rejected"):
            self.plan_fail_count += 1

    def _tracker_status_cb(self, msg):
        self.tracker_status = msg.data

    def _plan_cb(self, msg):
        if len(msg.poses) < 2:
            self.latest_path = []
            return
        self.latest_path = [
            (
                pose.pose.position.x,
                pose.pose.position.y,
                yaw_from_quaternion(pose.pose.orientation),
            )
            for pose in msg.poses
        ]
        self.plan_count += 1
        if self.pending_goal_time is not None:
            latency = (self.get_clock().now() - self.pending_goal_time).nanoseconds * 1e-9
            self.plan_latencies.append(max(0.0, latency))
            self.pending_goal_time = None
        self.plan_lengths.append(self._path_length(self.latest_path))
        self.plan_min_turning_radii.append(self._minimum_turning_radius(self.latest_path))
        self.reverse_fractions.append(self._reverse_fraction(self.latest_path))
        self.plan_collision_cells.append(self._occupied_path_points(self.latest_path))

    def _current_map_pose(self):
        if self.latest_odom is None:
            return None
        msg = self.latest_odom
        x = msg.pose.pose.position.x
        y = msg.pose.pose.position.y
        yaw = yaw_from_quaternion(msg.pose.pose.orientation)
        source_frame = msg.header.frame_id or self.map_frame
        if source_frame == self.map_frame:
            return x, y, yaw
        try:
            transform = self.tf_buffer.lookup_transform(
                self.map_frame,
                source_frame,
                Time(),
                timeout=Duration(seconds=self.tf_timeout),
            )
        except TransformException:
            self.transform_failures += 1
            return None
        t = transform.transform.translation
        transform_yaw = yaw_from_quaternion(transform.transform.rotation)
        cos_yaw = math.cos(transform_yaw)
        sin_yaw = math.sin(transform_yaw)
        return (
            t.x + cos_yaw * x - sin_yaw * y,
            t.y + sin_yaw * x + cos_yaw * y,
            normalize_angle(transform_yaw + yaw),
        )

    def _sample(self):
        pose = self._current_map_pose()
        if pose is None or self.latest_odom is None:
            return
        x, y, yaw = pose
        measured_linear = self.latest_odom.twist.twist.linear.x
        measured_angular = self.latest_odom.twist.twist.angular.z
        self.max_abs_measured_linear = max(
            self.max_abs_measured_linear, abs(measured_linear)
        )
        self.max_abs_measured_angular = max(
            self.max_abs_measured_angular, abs(measured_angular)
        )
        cross_track = self._cross_track_error(x, y)
        goal_error = None
        if self.latest_goal is not None:
            goal_error = math.hypot(self.latest_goal[0] - x, self.latest_goal[1] - y)
            self.goal_errors.append(goal_error)
        if cross_track is not None:
            self.cross_track_errors.append(cross_track)
        elapsed = (self.get_clock().now() - self.started_at).nanoseconds * 1e-9
        self.csv_writer.writerow([
            f"{elapsed:.6f}", f"{x:.6f}", f"{y:.6f}", f"{yaw:.6f}",
            f"{measured_linear:.6f}", f"{measured_angular:.6f}",
            f"{self.latest_cmd.linear.x:.6f}", f"{self.latest_cmd.angular.z:.6f}",
            f"{self.latest_drive.drive.speed:.6f}",
            f"{self.latest_drive.drive.steering_angle:.6f}",
            "" if cross_track is None else f"{cross_track:.6f}",
            "" if goal_error is None else f"{goal_error:.6f}",
            len(self.latest_path), self.planner_status, self.tracker_status,
        ])
        self.samples += 1
        if self.samples % 20 == 0:
            self.csv_file.flush()

    def _cross_track_error(self, x, y):
        if len(self.latest_path) < 2:
            return None
        return min(
            point_segment_distance(x, y, a[0], a[1], b[0], b[1])
            for a, b in zip(self.latest_path, self.latest_path[1:])
        )

    @staticmethod
    def _path_length(path):
        return sum(
            math.hypot(b[0] - a[0], b[1] - a[1])
            for a, b in zip(path, path[1:])
        )

    @staticmethod
    def _minimum_turning_radius(path):
        max_curvature = 0.0
        for a, b in zip(path, path[1:]):
            distance = math.hypot(b[0] - a[0], b[1] - a[1])
            if distance <= 1e-6:
                continue
            curvature = abs(normalize_angle(b[2] - a[2])) / distance
            max_curvature = max(max_curvature, curvature)
        return None if max_curvature <= 1e-9 else 1.0 / max_curvature

    @staticmethod
    def _reverse_fraction(path):
        reverse_length = 0.0
        total_length = 0.0
        for a, b in zip(path, path[1:]):
            dx = b[0] - a[0]
            dy = b[1] - a[1]
            distance = math.hypot(dx, dy)
            if distance <= 1e-6:
                continue
            segment_yaw = math.atan2(dy, dx)
            if math.cos(normalize_angle(a[2] - segment_yaw)) < 0.0:
                reverse_length += distance
            total_length += distance
        return 0.0 if total_length <= 1e-9 else reverse_length / total_length

    def _occupied_path_points(self, path):
        msg = self.map_msg
        if msg is None or msg.info.resolution <= 0.0:
            return None
        origin_yaw = yaw_from_quaternion(msg.info.origin.orientation)
        cos_yaw = math.cos(origin_yaw)
        sin_yaw = math.sin(origin_yaw)
        collisions = 0
        for x, y, _ in path:
            dx = x - msg.info.origin.position.x
            dy = y - msg.info.origin.position.y
            mx = cos_yaw * dx + sin_yaw * dy
            my = -sin_yaw * dx + cos_yaw * dy
            gx = math.floor(mx / msg.info.resolution)
            gy = math.floor(my / msg.info.resolution)
            if gx < 0 or gy < 0 or gx >= msg.info.width or gy >= msg.info.height:
                collisions += 1
                continue
            value = msg.data[gy * msg.info.width + gx]
            if value < 0 or value >= self.occupied_threshold:
                collisions += 1
        return collisions

    @staticmethod
    def _stats(values):
        if not values:
            return {"count": 0, "mean": None, "rms": None, "p95": None, "max": None}
        return {
            "count": len(values),
            "mean": statistics.fmean(values),
            "rms": math.sqrt(statistics.fmean(value * value for value in values)),
            "p95": percentile(values, 0.95),
            "max": max(values),
        }

    def write_summary(self):
        duration = (self.get_clock().now() - self.started_at).nanoseconds * 1e-9
        valid_radii = [value for value in self.plan_min_turning_radii if value is not None]
        valid_collisions = [value for value in self.plan_collision_cells if value is not None]
        summary = {
            "test_name": self.test_name,
            "duration_sec": duration,
            "samples": self.samples,
            "plan_count": self.plan_count,
            "plan_ok_status_count": self.plan_ok_count,
            "plan_fail_status_count": self.plan_fail_count,
            "planning_latency_sec": self._stats(self.plan_latencies),
            "cancel_to_zero_command_sec": self._stats(self.cancel_stop_latencies),
            "path_length_m": self._stats(self.plan_lengths),
            "path_occupied_points_max": max(valid_collisions) if valid_collisions else None,
            "minimum_planned_turning_radius_m": min(valid_radii) if valid_radii else None,
            "reverse_fraction": self._stats(self.reverse_fractions),
            "cross_track_error_m": self._stats(self.cross_track_errors),
            "goal_error_m": self._stats(self.goal_errors),
            "final_goal_error_m": self.goal_errors[-1] if self.goal_errors else None,
            "minimum_goal_error_m": min(self.goal_errors) if self.goal_errors else None,
            "transform_failures": self.transform_failures,
            "max_abs_cmd_linear": self.max_abs_cmd_linear,
            "max_abs_cmd_angular": self.max_abs_cmd_angular,
            "max_abs_drive_speed": self.max_abs_drive_speed,
            "max_abs_steering_angle": self.max_abs_steering_angle,
            "max_abs_measured_linear": self.max_abs_measured_linear,
            "max_abs_measured_angular": self.max_abs_measured_angular,
            "last_planner_status": self.planner_status,
            "last_tracker_status": self.tracker_status,
            "csv": str(self.csv_path),
        }
        self.csv_file.flush()
        self.summary_path.write_text(
            json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8"
        )
        self.get_logger().info(f"wrote navigation summary to {self.summary_path}")

    def close(self):
        if not self.csv_file.closed:
            self.write_summary()
            self.csv_file.close()


def main():
    rclpy.init()
    node = NavigationTestRecorder()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.close()
        try:
            if rclpy.ok():
                node.destroy_node()
        except Exception:
            pass
        try:
            if rclpy.ok():
                rclpy.shutdown()
        except Exception:
            pass


if __name__ == "__main__":
    main()
