#!/usr/bin/env python3
"""Velocity smoother node with low-pass Butterworth filter.

Subscribes to /odom, applies low-pass Butterworth filter to linear and angular
velocity components, and publishes smoothed odometry on /odom_smooth.

Reduces noise from encoder jitter and improves state estimation stability.
"""

import math
import rclpy
from rclpy.node import Node
from nav_msgs.msg import Odometry
from std_msgs.msg import Float32


class ButterworthFilter:
    """Simple second-order Butterworth low-pass filter for continuous signals."""

    def __init__(self, cutoff_hz, sample_rate_hz, order=2):
        """Initialize Butterworth filter.

        Args:
            cutoff_hz: Cutoff frequency in Hz
            sample_rate_hz: Sampling rate in Hz
            order: Filter order (typically 1-4)
        """
        self.cutoff_hz = cutoff_hz
        self.sample_rate_hz = sample_rate_hz
        self.order = order

        # Normalized frequency (0 to 1, where 1 = Nyquist)
        self.omega_n = 2.0 * math.pi * cutoff_hz / sample_rate_hz

        # Simplified filter coefficients for order 2
        # Using canonical form: y[n] = b0*x[n] + b1*x[n-1] + b2*x[n-2] - a1*y[n-1] - a2*y[n-2]
        if order == 1:
            # First-order filter
            alpha = self.omega_n / (self.omega_n + 2.0)
            self.b = [alpha, alpha]
            self.a = [1.0, -(1.0 - 2.0 * alpha)]
        else:
            # Second-order filter (butterworth)
            sqrt2 = math.sqrt(2.0)
            wc = math.tan(self.omega_n / 2.0)
            wc2 = wc * wc

            denom = 1.0 + sqrt2 * wc + wc2

            self.b = [wc2 / denom, 2.0 * wc2 / denom, wc2 / denom]
            self.a = [1.0,
                      2.0 * (wc2 - 1.0) / denom,
                      (1.0 - sqrt2 * wc + wc2) / denom]

        # State buffers
        self.x_history = [0.0, 0.0, 0.0]  # Input history
        self.y_history = [0.0, 0.0]       # Output history

    def filter(self, x):
        """Apply filter to input value.

        Args:
            x: Input value

        Returns:
            Filtered output value
        """
        # Update input history
        self.x_history[2] = self.x_history[1]
        self.x_history[1] = self.x_history[0]
        self.x_history[0] = x

        # Compute output using difference equation
        if len(self.b) == 2:
            # First-order filter
            y = (self.b[0] * self.x_history[0] +
                 self.b[1] * self.x_history[1] -
                 self.a[1] * self.y_history[1])
        else:
            # Second-order filter
            y = (self.b[0] * self.x_history[0] +
                 self.b[1] * self.x_history[1] +
                 self.b[2] * self.x_history[2] -
                 self.a[1] * self.y_history[1] -
                 self.a[2] * self.y_history[2])

        # Update output history
        self.y_history[1] = self.y_history[0]
        self.y_history[0] = y

        return y

    def reset(self):
        """Reset filter state."""
        self.x_history = [0.0, 0.0, 0.0]
        self.y_history = [0.0, 0.0]


class VelocitySmootherNode(Node):
    """ROS2 node for velocity smoothing via low-pass filtering."""

    def __init__(self):
        super().__init__('velocity_smoother')

        # Parameters
        self.declare_parameter('filter_order', 2)
        self.declare_parameter('cutoff_frequency', 5.0)
        self.declare_parameter('publish_frequency', 50)
        self.declare_parameter('enable_smoothing', True)

        filter_order = self.get_parameter('filter_order').value
        cutoff_frequency = self.get_parameter('cutoff_frequency').value
        publish_frequency = self.get_parameter('publish_frequency').value
        self.enable_smoothing = self.get_parameter('enable_smoothing').value

        # Create filters for each velocity component
        self.filter_linear_x = ButterworthFilter(
            cutoff_frequency, publish_frequency, order=filter_order
        )
        self.filter_linear_y = ButterworthFilter(
            cutoff_frequency, publish_frequency, order=filter_order
        )
        self.filter_linear_z = ButterworthFilter(
            cutoff_frequency, publish_frequency, order=filter_order
        )
        self.filter_angular_x = ButterworthFilter(
            cutoff_frequency, publish_frequency, order=filter_order
        )
        self.filter_angular_y = ButterworthFilter(
            cutoff_frequency, publish_frequency, order=filter_order
        )
        self.filter_angular_z = ButterworthFilter(
            cutoff_frequency, publish_frequency, order=filter_order
        )

        # Last received odometry
        self.last_odom = None

        # Subscription to raw odometry
        self.sub_odom = self.create_subscription(
            Odometry, '/odom', self._on_odom, 10
        )

        # Publisher for smoothed odometry
        self.pub_odom_smooth = self.create_publisher(Odometry, '/odom_smooth', 10)

        # Timer for publishing at fixed frequency
        period = 1.0 / publish_frequency
        self.timer = self.create_timer(period, self._timer_callback)

        self.get_logger().info(
            f"Velocity smoother initialized. "
            f"Cutoff: {cutoff_frequency}Hz, Order: {filter_order}, "
            f"Publish: {publish_frequency}Hz"
        )

    def _on_odom(self, msg: Odometry) -> None:
        """Callback for incoming odometry messages."""
        self.last_odom = msg

    def _timer_callback(self) -> None:
        """Periodically filter and publish smoothed odometry."""
        if self.last_odom is None:
            return

        # Create output message
        smoothed = Odometry()
        smoothed.header = self.last_odom.header
        smoothed.child_frame_id = self.last_odom.child_frame_id

        # Copy pose (unchanged)
        smoothed.pose = self.last_odom.pose

        if self.enable_smoothing:
            # Filter velocity components
            linear_x = self.filter_linear_x.filter(
                self.last_odom.twist.twist.linear.x
            )
            linear_y = self.filter_linear_y.filter(
                self.last_odom.twist.twist.linear.y
            )
            linear_z = self.filter_linear_z.filter(
                self.last_odom.twist.twist.linear.z
            )

            angular_x = self.filter_angular_x.filter(
                self.last_odom.twist.twist.angular.x
            )
            angular_y = self.filter_angular_y.filter(
                self.last_odom.twist.twist.angular.y
            )
            angular_z = self.filter_angular_z.filter(
                self.last_odom.twist.twist.angular.z
            )
        else:
            # No filtering
            linear_x = self.last_odom.twist.twist.linear.x
            linear_y = self.last_odom.twist.twist.linear.y
            linear_z = self.last_odom.twist.twist.linear.z
            angular_x = self.last_odom.twist.twist.angular.x
            angular_y = self.last_odom.twist.twist.angular.y
            angular_z = self.last_odom.twist.twist.angular.z

        # Set smoothed twist
        smoothed.twist.twist.linear.x = linear_x
        smoothed.twist.twist.linear.y = linear_y
        smoothed.twist.twist.linear.z = linear_z
        smoothed.twist.twist.angular.x = angular_x
        smoothed.twist.twist.angular.y = angular_y
        smoothed.twist.twist.angular.z = angular_z

        # Copy covariances
        smoothed.twist.covariance = self.last_odom.twist.covariance

        self.pub_odom_smooth.publish(smoothed)


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


if __name__ == '__main__':
    main()