#!/usr/bin/env python3
from typing import Optional
import math
import rclpy
from rclpy.node import Node
from rclpy.timer import Timer
from geometry_msgs.msg import Twist
from nav_msgs.msg import Odometry
from std_msgs.msg import Bool

class WheelSlipDetectorNode(Node):
    def __init__(self):
        super().__init__("wheel_slip_detector")
        self.declare_parameter("frequency", 10)
        frequency = self.get_parameter("frequency").value
        self.declare_parameter("slip_threshold", 0.1)
        self.declare_parameter("slip_timeout", 0.5)
        self.slip_threshold = self.get_parameter("slip_threshold").value
        self.slip_timeout = self.get_parameter("slip_timeout").value
        self.period = 1.0 / frequency
        self.cmd_vel: Optional[Twist] = None
        self.actual_vel: Optional[Twist] = None
        self.slip_duration = 0.0
        self.slip_detected = False
        self.create_subscription(Twist, "/cmd_vel", self._on_cmd_vel, 10)
        self.create_subscription(Odometry, "/odom", self._on_odom, 10)
        self.pub_slip = self.create_publisher(Bool, "/saltybot/wheel_slip_detected", 10)
        self.timer: Timer = self.create_timer(self.period, self._timer_callback)
        self.get_logger().info(f"Wheel slip detector initialized at {frequency}Hz. Threshold: {self.slip_threshold} m/s, Timeout: {self.slip_timeout}s")

    def _on_cmd_vel(self, msg: Twist) -> None:
        self.cmd_vel = msg

    def _on_odom(self, msg: Odometry) -> None:
        self.actual_vel = msg.twist.twist

    def _timer_callback(self) -> None:
        if self.cmd_vel is None or self.actual_vel is None:
            slip_detected = False
        else:
            slip_detected = self._check_slip()
        if slip_detected:
            self.slip_duration += self.period
        else:
            self.slip_duration = 0.0
        is_slip = self.slip_duration > self.slip_timeout
        if is_slip != self.slip_detected:
            self.slip_detected = is_slip
            if self.slip_detected:
                self.get_logger().warn(f"WHEEL SLIP DETECTED: {self.slip_duration:.2f}s")
            else:
                self.get_logger().info("Wheel slip cleared")
        slip_msg = Bool()
        slip_msg.data = is_slip
        self.pub_slip.publish(slip_msg)

    def _check_slip(self) -> bool:
        cmd_speed = math.sqrt(self.cmd_vel.linear.x**2 + self.cmd_vel.linear.y**2)
        actual_speed = math.sqrt(self.actual_vel.linear.x**2 + self.actual_vel.linear.y**2)
        vel_diff = abs(cmd_speed - actual_speed)
        if cmd_speed < 0.05 and actual_speed < 0.05:
            return False
        return vel_diff > self.slip_threshold

def main(args=None):
    rclpy.init(args=args)
    node = WheelSlipDetectorNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == "__main__":
    main()
