#!/usr/bin/env python3
"""
Smooth velocity controller with accel/decel ramp.

Subscribes to /cmd_vel (geometry_msgs/Twist), applies acceleration
and deceleration ramps, and publishes smoothed commands to
/cmd_vel_smoothed at 50 Hz.

E-stop (std_msgs/Bool on /e_stop): immediate zero, bypasses ramp.
Optional jerk limiting via max_linear_jerk / max_angular_jerk params.
"""

import math

import rclpy
from geometry_msgs.msg import Twist
from rclpy.node import Node
from std_msgs.msg import Bool


class VelocitySmoother(Node):
    def __init__(self):
        super().__init__('velocity_smoother')

        # Declare parameters
        self.declare_parameter('publish_rate', 50.0)           # Hz
        self.declare_parameter('max_linear_accel', 1.0)        # m/s²
        self.declare_parameter('max_linear_decel', 2.0)        # m/s²
        self.declare_parameter('max_angular_accel', 1.5)       # rad/s²
        self.declare_parameter('max_angular_decel', 3.0)       # rad/s²
        self.declare_parameter('max_linear_jerk', 0.0)         # m/s³, 0=disabled
        self.declare_parameter('max_angular_jerk', 0.0)        # rad/s³, 0=disabled

        rate = self.get_parameter('publish_rate').value
        self.max_lin_acc = self.get_parameter('max_linear_accel').value
        self.max_lin_dec = self.get_parameter('max_linear_decel').value
        self.max_ang_acc = self.get_parameter('max_angular_accel').value
        self.max_ang_dec = self.get_parameter('max_angular_decel').value
        self.max_lin_jerk = self.get_parameter('max_linear_jerk').value
        self.max_ang_jerk = self.get_parameter('max_angular_jerk').value

        self._dt = 1.0 / rate

        # State
        self._target_lin = 0.0
        self._target_ang = 0.0
        self._current_lin = 0.0
        self._current_ang = 0.0
        self._current_lin_acc = 0.0   # for jerk limiting
        self._current_ang_acc = 0.0
        self._e_stop = False

        # Publisher
        self._pub = self.create_publisher(Twist, '/cmd_vel_smoothed', 10)

        # Subscriptions
        self.create_subscription(Twist, '/cmd_vel', self._cmd_vel_cb, 10)
        self.create_subscription(Bool, '/e_stop', self._e_stop_cb, 10)

        # Timer
        self.create_timer(self._dt, self._timer_cb)

        self.get_logger().info(
            f'VelocitySmoother ready at {rate:.0f} Hz — '
            f'lin_acc={self.max_lin_acc} lin_dec={self.max_lin_dec} '
            f'ang_acc={self.max_ang_acc} ang_dec={self.max_ang_dec}'
        )

    # ------------------------------------------------------------------

    def _cmd_vel_cb(self, msg: Twist):
        self._target_lin = msg.linear.x
        self._target_ang = msg.angular.z

    def _e_stop_cb(self, msg: Bool):
        self._e_stop = msg.data
        if self._e_stop:
            self._target_lin = 0.0
            self._target_ang = 0.0
            self._current_lin = 0.0
            self._current_ang = 0.0
            self._current_lin_acc = 0.0
            self._current_ang_acc = 0.0
            self.get_logger().warn('E-STOP active — motors zeroed immediately')

    # ------------------------------------------------------------------

    def _ramp(
        self,
        current: float,
        target: float,
        max_acc: float,
        max_dec: float,
        current_acc: float,
        max_jerk: float,
    ) -> tuple[float, float]:
        """
        Advance `current` toward `target` with separate accel/decel limits.
        Optionally apply jerk limiting to the acceleration.
        Returns (new_value, new_acc).
        """
        error = target - current

        # Choose limit: decelerate if moving away from zero and target is
        # closer to zero (or past it), else accelerate.
        if current * error < 0 or (error != 0 and abs(target) < abs(current)):
            limit = max_dec
        else:
            limit = max_acc

        desired_acc = math.copysign(min(abs(error) / self._dt, limit), error)
        # Clamp so we don't overshoot
        desired_acc = max(-limit, min(limit, desired_acc))

        if max_jerk > 0.0:
            max_d_acc = max_jerk * self._dt
            new_acc = current_acc + max(
                -max_d_acc, min(max_d_acc, desired_acc - current_acc)
            )
        else:
            new_acc = desired_acc

        delta = new_acc * self._dt
        # Clamp so we don't overshoot the target
        if abs(delta) > abs(error):
            delta = error
        return current + delta, new_acc

    def _timer_cb(self):
        if self._e_stop:
            msg = Twist()
            self._pub.publish(msg)
            return

        self._current_lin, self._current_lin_acc = self._ramp(
            self._current_lin,
            self._target_lin,
            self.max_lin_acc,
            self.max_lin_dec,
            self._current_lin_acc,
            self.max_lin_jerk,
        )
        self._current_ang, self._current_ang_acc = self._ramp(
            self._current_ang,
            self._target_ang,
            self.max_ang_acc,
            self.max_ang_dec,
            self._current_ang_acc,
            self.max_lin_jerk,  # intentional: use linear jerk scale for angular too if set
        )

        msg = Twist()
        msg.linear.x = self._current_lin
        msg.angular.z = self._current_ang
        self._pub.publish(msg)


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


if __name__ == '__main__':
    main()
