#!/usr/bin/env python3
"""Smooth velocity controller - Issue #455. Acceleration limiting + S-curve jerk reduction."""

import math
import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Twist
from std_msgs.msg import String

class SmoothVelocityNode(Node):
    def __init__(self):
        super().__init__("smooth_velocity_controller")
        self.declare_parameter("max_linear_accel", 0.5)
        self.declare_parameter("max_angular_accel", 1.0)
        self.declare_parameter("max_linear_decel", 0.8)
        self.declare_parameter("max_angular_decel", 1.0)
        self.declare_parameter("use_scurve", True)
        self.declare_parameter("scurve_duration", 0.2)
        self.declare_parameter("update_rate", 50)

        self.max_lin_accel = self.get_parameter("max_linear_accel").value
        self.max_ang_accel = self.get_parameter("max_angular_accel").value
        self.max_lin_decel = self.get_parameter("max_linear_decel").value
        self.max_ang_decel = self.get_parameter("max_angular_decel").value
        self.use_scurve = self.get_parameter("use_scurve").value
        self.scurve_duration = self.get_parameter("scurve_duration").value
        self.update_rate = self.get_parameter("update_rate").value
        self.dt = 1.0 / self.update_rate

        self.cmd_vel_pub = self.create_publisher(Twist, "/cmd_vel", 1)
        self.velocity_profile_pub = self.create_publisher(String, "/saltybot/velocity_profile", 1)
        self.create_subscription(Twist, "/cmd_vel_raw", self._cmd_vel_raw_callback, 1)

        self.current_lin_vel = 0.0
        self.current_ang_vel = 0.0
        self.target_lin_vel = 0.0
        self.target_ang_vel = 0.0
        self.scurve_time_lin = 0.0
        self.scurve_time_ang = 0.0

        self.create_timer(self.dt, self._update_loop)
        self.get_logger().info(f"Smooth velocity controller: {self.update_rate}Hz, S-curve={self.use_scurve}")

    def _cmd_vel_raw_callback(self, msg: Twist):
        self.target_lin_vel = msg.linear.x
        self.target_ang_vel = msg.angular.z

    def _update_loop(self):
        if self.use_scurve:
            lin_vel = self._apply_scurve(self.current_lin_vel, self.target_lin_vel, self.max_lin_accel, self.max_lin_decel, "linear")
            ang_vel = self._apply_scurve(self.current_ang_vel, self.target_ang_vel, self.max_ang_accel, self.max_ang_decel, "angular")
        else:
            lin_vel = self._apply_accel_limit(self.current_lin_vel, self.target_lin_vel, self.max_lin_accel, self.max_lin_decel)
            ang_vel = self._apply_accel_limit(self.current_ang_vel, self.target_ang_vel, self.max_ang_accel, self.max_ang_decel)

        self.current_lin_vel = lin_vel
        self.current_ang_vel = ang_vel

        twist = Twist()
        twist.linear.x = lin_vel
        twist.angular.z = ang_vel
        self.cmd_vel_pub.publish(twist)

        profile = f"lin:{lin_vel:.3f},ang:{ang_vel:.3f}"
        self.velocity_profile_pub.publish(String(data=profile))

    def _apply_scurve(self, current, target, max_accel, max_decel, axis):
        if abs(target - current) < 0.001:
            return target
        direction = 1 if target > current else -1
        accel_limit = max_accel if direction > 0 else max_decel
        max_vel_change = accel_limit * self.dt
        if self.scurve_duration > 0:
            time_var = self.scurve_time_lin if axis == "linear" else self.scurve_time_ang
            progress = min(time_var, self.scurve_duration) / self.scurve_duration
            smooth_factor = (1 - math.cos(progress * math.pi)) / 2
            max_vel_change *= smooth_factor
        new_vel = current + max_vel_change * direction
        if (direction > 0 and new_vel >= target) or (direction < 0 and new_vel <= target):
            if axis == "linear":
                self.scurve_time_lin = 0.0
            else:
                self.scurve_time_ang = 0.0
            return target
        if axis == "linear":
            self.scurve_time_lin += self.dt
        else:
            self.scurve_time_ang += self.dt
        return new_vel

    def _apply_accel_limit(self, current, target, max_accel, max_decel):
        if abs(target - current) < 0.001:
            return target
        direction = 1 if target > current else -1
        accel_limit = max_accel if direction > 0 else max_decel
        max_vel_change = accel_limit * self.dt
        new_vel = current + max_vel_change * direction
        if (direction > 0 and new_vel >= target) or (direction < 0 and new_vel <= target):
            return target
        return new_vel

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

if __name__ == "__main__":
    main()
