#!/usr/bin/env python3
"""Unit tests for VelocitySmoother ramp logic."""

import math
import unittest
from unittest.mock import MagicMock, patch


# ---------------------------------------------------------------------------
# Minimal ROS2 stubs so tests run without a ROS2 installation
# ---------------------------------------------------------------------------

class _FakeParam:
    def __init__(self, val):
        self._val = val
    @property
    def value(self):
        return self._val


class _FakeNode:
    def __init__(self):
        self._params = {}
    def declare_parameter(self, name, default):
        self._params[name] = _FakeParam(default)
    def get_parameter(self, name):
        return self._params[name]
    def create_publisher(self, *a, **kw):
        return MagicMock()
    def create_subscription(self, *a, **kw):
        return MagicMock()
    def create_timer(self, *a, **kw):
        return MagicMock()
    def get_logger(self):
        log = MagicMock()
        log.info = MagicMock()
        log.warn = MagicMock()
        return log


# Patch rclpy and geometry_msgs before importing the module.
# rclpy.node.Node must be a *real* class so that VelocitySmoother can
# inherit from it without becoming a MagicMock itself.
import sys

class _RealNodeBase:
    """Minimal real base class that stands in for rclpy.node.Node."""
    def __init__(self, *args, **kwargs):
        pass

    def declare_parameter(self, name, default):
        if not hasattr(self, '_params'):
            self._params = {}
        self._params[name] = _FakeParam(default)

    def get_parameter(self, name):
        return self._params[name]

    def create_publisher(self, *a, **kw):
        return MagicMock()

    def create_subscription(self, *a, **kw):
        return MagicMock()

    def create_timer(self, *a, **kw):
        return MagicMock()

    def get_logger(self):
        log = MagicMock()
        log.info = MagicMock()
        log.warn = MagicMock()
        return log

rclpy_node_mod = MagicMock()
rclpy_node_mod.Node = _RealNodeBase

rclpy_mock = MagicMock()
rclpy_mock.node = rclpy_node_mod

sys.modules['rclpy'] = rclpy_mock
sys.modules['rclpy.node'] = rclpy_node_mod
sys.modules.setdefault('geometry_msgs', MagicMock())
sys.modules.setdefault('geometry_msgs.msg', MagicMock())
sys.modules.setdefault('std_msgs', MagicMock())
sys.modules.setdefault('std_msgs.msg', MagicMock())

# Provide a real Twist-like object for tests
class _Twist:
    class _Vec:
        x = 0.0
        y = 0.0
        z = 0.0
    def __init__(self):
        self.linear = self._Vec()
        self.angular = self._Vec()

sys.modules['geometry_msgs.msg'].Twist = _Twist
sys.modules['std_msgs.msg'].Bool = MagicMock()

import importlib, os, pathlib
sys.path.insert(0, str(pathlib.Path(__file__).parents[1] / 'saltybot_vesc_driver'))

# Now we can import the smoother logic directly
from velocity_smoother_node import VelocitySmoother


def _make_smoother(**params):
    """Create a VelocitySmoother backed by _RealNodeBase with custom params."""
    defaults = dict(
        publish_rate=50.0,
        max_linear_accel=1.0,
        max_linear_decel=2.0,
        max_angular_accel=1.5,
        max_angular_decel=3.0,
        max_linear_jerk=0.0,
        max_angular_jerk=0.0,
    )
    defaults.update(params)

    # Pre-seed _params before __init__ so declare_parameter picks up overrides.
    node = VelocitySmoother.__new__(VelocitySmoother)
    node._params = {k: _FakeParam(v) for k, v in defaults.items()}

    # Monkey-patch declare_parameter so it doesn't overwrite our pre-seeded values
    original_declare = _RealNodeBase.declare_parameter
    def _noop_declare(self, name, default):
        pass  # params already seeded
    _RealNodeBase.declare_parameter = _noop_declare

    try:
        VelocitySmoother.__init__(node)
    finally:
        _RealNodeBase.declare_parameter = original_declare

    return node


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

class TestRampLogic(unittest.TestCase):

    def _make(self, **kw):
        return _make_smoother(**kw)

    def test_ramp_reaches_target_within_expected_steps(self):
        """From 0 to 1 m/s with accel=1 m/s² at 50 Hz → ~50 steps."""
        node = self._make(max_linear_accel=1.0, publish_rate=50.0)
        node._target_lin = 1.0
        steps = 0
        while abs(node._current_lin - 1.0) > 0.01 and steps < 200:
            node._timer_cb()
            steps += 1
        self.assertLessEqual(steps, 55, "Should reach 1 m/s within ~55 steps at 50 Hz")
        self.assertAlmostEqual(node._current_lin, 1.0, places=2)

    def test_decel_faster_than_accel(self):
        """Deceleration should complete faster than acceleration."""
        node = self._make(max_linear_accel=1.0, max_linear_decel=2.0, publish_rate=50.0)
        # Ramp up
        node._target_lin = 1.0
        for _ in range(100):
            node._timer_cb()
        # Now decelerate
        node._target_lin = 0.0
        decel_steps = 0
        while abs(node._current_lin) > 0.01 and decel_steps < 200:
            node._timer_cb()
            decel_steps += 1
        # Ramp up again to measure accel steps
        node._current_lin = 0.0
        node._target_lin = 1.0
        accel_steps = 0
        while abs(node._current_lin - 1.0) > 0.01 and accel_steps < 200:
            node._timer_cb()
            accel_steps += 1
        self.assertLess(decel_steps, accel_steps,
                        "Decel (2 m/s²) should finish in fewer steps than accel (1 m/s²)")

    def test_e_stop_zeros_immediately(self):
        """E-stop must zero velocity in the same callback, bypassing ramp."""
        node = self._make()
        node._current_lin = 2.0
        node._current_ang = 1.0
        node._target_lin = 2.0
        node._target_ang = 1.0

        msg = MagicMock()
        msg.data = True
        node._e_stop_cb(msg)

        self.assertEqual(node._current_lin, 0.0)
        self.assertEqual(node._current_ang, 0.0)
        self.assertTrue(node._e_stop)

    def test_no_overshoot(self):
        """Current velocity must never exceed target during ramp-up."""
        node = self._make(max_linear_accel=1.0, publish_rate=50.0)
        node._target_lin = 0.5
        for _ in range(100):
            node._timer_cb()
            self.assertLessEqual(node._current_lin, 0.5 + 1e-9)

    def test_negative_velocity_ramp(self):
        """Ramp works symmetrically for negative targets."""
        node = self._make(max_linear_accel=1.0, publish_rate=50.0)
        node._target_lin = -1.0
        for _ in range(200):
            node._timer_cb()
        self.assertAlmostEqual(node._current_lin, -1.0, places=2)

    def test_angular_ramp(self):
        """Angular velocity ramps correctly."""
        node = self._make(max_angular_accel=1.5, publish_rate=50.0)
        node._target_ang = 1.0
        for _ in range(200):
            node._timer_cb()
        self.assertAlmostEqual(node._current_ang, 1.0, places=2)

    def test_ramp_timing_linear(self):
        """Time to ramp from 0 to v_max ≈ v_max / accel (±10%)."""
        accel = 1.0  # m/s²
        v_max = 1.0  # m/s
        rate = 50.0
        expected_time = v_max / accel  # 1.0 s
        node = self._make(max_linear_accel=accel, publish_rate=rate)
        node._target_lin = v_max
        steps = 0
        while abs(node._current_lin - v_max) > 0.01 and steps < 500:
            node._timer_cb()
            steps += 1
        actual_time = steps / rate
        self.assertAlmostEqual(actual_time, expected_time, delta=expected_time * 0.12,
                               msg=f"Ramp time {actual_time:.2f}s expected ~{expected_time:.2f}s")


if __name__ == '__main__':
    unittest.main()
