"""Unit tests for pid_autotune_node."""

import pytest
import json
import math
from geometry_msgs.msg import Twist
from sensor_msgs.msg import Imu
from std_msgs.msg import Float32, String

import rclpy

# Import the node under test
from saltybot_pid_autotune.pid_autotune_node import (
    PIDAutotuneNode,
    TuneState,
)


@pytest.fixture
def rclpy_fixture():
    """Initialize and cleanup rclpy."""
    rclpy.init()
    yield
    rclpy.shutdown()


@pytest.fixture
def node(rclpy_fixture):
    """Create a PID autotune node instance."""
    node = PIDAutotuneNode()
    yield node
    node.destroy_node()


class TestPIDAutotuneNode:
    """Test suite for PIDAutotuneNode."""

    def test_node_initialization(self, node):
        """Test that node initializes with correct defaults."""
        assert node.tune_state == TuneState.IDLE
        assert node.relay_magnitude == 0.3
        assert node.tilt_safety_limit == 25.0
        assert node.max_tune_duration == 120.0
        assert node.oscillation_count == 5
        assert node.last_imu_tilt == 0.0
        assert node.last_feedback == 0.0
        assert len(node.feedback_history) == 0
        assert len(node.peaks) == 0

    def test_imu_tilt_extraction(self, node):
        """Test IMU tilt angle extraction."""
        msg = Imu()
        # ax=1, ay=0, az=0 -> pitch = atan2(1, 0) = 90 degrees
        msg.linear_acceleration.x = 1.0
        msg.linear_acceleration.y = 0.0
        msg.linear_acceleration.z = 0.0

        node._on_imu(msg)
        assert abs(node.last_imu_tilt - 90.0) < 1.0

    def test_imu_small_tilt(self, node):
        """Test small tilt angle extraction."""
        msg = Imu()
        # Simulate 10 degree tilt
        accel_x = math.sin(math.radians(10))
        accel_z = math.cos(math.radians(10))
        msg.linear_acceleration.x = accel_x
        msg.linear_acceleration.y = 0.0
        msg.linear_acceleration.z = accel_z

        node._on_imu(msg)
        assert abs(node.last_imu_tilt - 10.0) < 1.0

    def test_feedback_recording(self, node):
        """Test feedback value recording."""
        msg = Float32(data=0.5)
        node._on_feedback(msg)
        assert node.last_feedback == 0.5

    def test_feedback_history_during_tuning(self, node):
        """Test that feedback history is only recorded during MEASURING state."""
        # Not tuning yet
        node._on_feedback(Float32(data=1.0))
        assert len(node.feedback_history) == 0

        # Start measuring
        node.tune_state = TuneState.MEASURING
        node.start_time = node.get_clock().now().nanoseconds / 1e9

        node._on_feedback(Float32(data=1.0))
        assert len(node.feedback_history) == 1

        node._on_feedback(Float32(data=2.0))
        assert len(node.feedback_history) == 2

    def test_ziegler_nichols_pd_gains(self, node):
        """Test Ziegler-Nichols PD gain calculation."""
        # Example: Ku=2.0, Tu=1.0
        kp, kd = node._compute_ziegler_nichols(2.0, 1.0)

        # Kp = 0.6 * Ku = 0.6 * 2.0 = 1.2
        assert abs(kp - 1.2) < 1e-6

        # Kd = 0.075 * Ku * Tu = 0.075 * 2.0 * 1.0 = 0.15
        assert abs(kd - 0.15) < 1e-6

    def test_ziegler_nichols_high_gains(self, node):
        """Test Z-N gains with higher Ku/Tu values."""
        kp, kd = node._compute_ziegler_nichols(3.0, 2.0)

        assert abs(kp - 1.8) < 1e-6  # 0.6 * 3.0
        assert abs(kd - 0.45) < 1e-6  # 0.075 * 3.0 * 2.0

    def test_ultimate_gain_from_relay(self, node):
        """Test Ku calculation from relay magnitude and peak."""
        # Ku = 4*relay_magnitude / (pi * |peak|)
        relay_mag = 0.3
        peak_amplitude = 0.5

        ku = (4.0 * relay_mag) / (math.pi * peak_amplitude)

        expected = (4.0 * 0.3) / (math.pi * 0.5)
        assert abs(ku - expected) < 1e-6

    def test_analyze_oscillation_insufficient_peaks(self, node):
        """Test that analysis fails with insufficient peaks."""
        node.peaks = [(0.0, 0.5), (1.0, -0.5)]  # Only 2 peaks

        ku, tu = node._analyze_oscillation()
        assert ku is None
        assert tu is None

    def test_analyze_oscillation_valid_data(self, node):
        """Test oscillation analysis with valid peak data."""
        node.relay_magnitude = 0.2
        # Simulate 5 alternating peaks every ~1 second
        node.peaks = [
            (0.0, -0.4),  # Negative peak
            (1.0, 0.4),   # Positive peak
            (2.0, -0.4),  # Negative peak
            (3.0, 0.4),   # Positive peak
            (4.0, -0.4),  # Negative peak
        ]

        ku, tu = node._analyze_oscillation()

        # Should compute meaningful values
        assert ku is not None
        assert tu is not None
        assert ku > 0
        assert tu > 0

    def test_tune_state_machine(self, node):
        """Test state machine transitions."""
        assert node.tune_state == TuneState.IDLE

        # Invalid state for tuning (should reject)
        node.tune_state = TuneState.TUNING
        response = node._handle_autotune_request(
            type("Request", (), {}), type("Response", (), {})
        )
        # Should fail because already tuning
        # (actual test depends on response structure)

    def test_info_json_format(self, node):
        """Test that info JSON is valid."""
        info = {
            "status": "measuring",
            "cycles": 3,
            "feedback": 0.5,
            "tilt": 5.0,
        }

        json_str = json.dumps(info)
        parsed = json.loads(json_str)

        assert parsed["status"] == "measuring"
        assert parsed["cycles"] == 3
        assert abs(parsed["feedback"] - 0.5) < 1e-6
        assert abs(parsed["tilt"] - 5.0) < 1e-6

    def test_tilt_safety_limit_check(self, node):
        """Test that tilt exceeding limit is detected."""
        node.last_imu_tilt = 30.0  # Exceeds 25 degree limit
        node.tilt_safety_limit = 25.0

        # In actual settle/relay code, this would trigger abort
        assert abs(node.last_imu_tilt) > node.tilt_safety_limit

    def test_tilt_safe_margin(self, node):
        """Test normal tilt well within limits."""
        node.last_imu_tilt = 10.0
        node.tilt_safety_limit = 25.0

        assert abs(node.last_imu_tilt) < node.tilt_safety_limit


class TestPIDAutotuneScenarios:
    """Integration-style tests for realistic scenarios."""

    def test_scenario_clean_oscillation(self, node):
        """Scenario: perfect sinusoidal oscillation."""
        node.relay_magnitude = 0.25

        # Simulate clean oscillation: 5 peaks, 1Hz frequency
        node.peaks = [
            (0.0, -0.5),
            (1.0, 0.5),
            (2.0, -0.5),
            (3.0, 0.5),
            (4.0, -0.5),
        ]

        ku, tu = node._analyze_oscillation()

        assert ku is not None
        assert tu is not None
        # Tu should be ~2 seconds (period of oscillation)
        assert 1.8 < tu < 2.2

    def test_scenario_damped_oscillation(self, node):
        """Scenario: damped oscillation with decaying peaks."""
        node.relay_magnitude = 0.2

        # Peaks decay over time
        node.peaks = [
            (0.0, -0.6),
            (1.0, 0.55),
            (2.0, -0.5),
            (3.0, 0.45),
            (4.0, -0.4),
        ]

        ku, tu = node._analyze_oscillation()

        # Should still work with average of peaks
        assert ku is not None
        assert tu is not None

    def test_scenario_high_frequency_oscillation(self, node):
        """Scenario: fast oscillation (0.2Hz)."""
        node.relay_magnitude = 0.3

        # 5 second period = 0.2 Hz
        node.peaks = [
            (0.0, -0.5),
            (5.0, 0.5),
            (10.0, -0.5),
            (15.0, 0.5),
            (20.0, -0.5),
        ]

        ku, tu = node._analyze_oscillation()

        assert tu is not None
        # Period should be ~10 seconds (5 second half-period)
        assert 9.8 < tu < 10.2

    def test_scenario_z_n_aggressive_tuning(self, node):
        """Scenario: Z-N gains for aggressive tuning."""
        # High ultimate gain
        kp, kd = node._compute_ziegler_nichols(ku=4.0, tu=0.5)

        assert kp == pytest.approx(2.4)  # 0.6 * 4.0
        assert kd == pytest.approx(0.15)  # 0.075 * 4.0 * 0.5

    def test_scenario_z_n_conservative_tuning(self, node):
        """Scenario: Z-N gains for conservative tuning."""
        # Low ultimate gain
        kp, kd = node._compute_ziegler_nichols(ku=1.0, tu=2.0)

        assert kp == pytest.approx(0.6)  # 0.6 * 1.0
        assert kd == pytest.approx(0.15)  # 0.075 * 1.0 * 2.0

    def test_scenario_tilt_safety_aborts(self, node):
        """Scenario: tuning aborted due to tilt exceeding limit."""
        node.last_imu_tilt = 28.0  # Exceeds 25 degree limit
        node.tilt_safety_limit = 25.0

        # Safety check would catch this
        should_abort = abs(node.last_imu_tilt) > node.tilt_safety_limit
        assert should_abort is True

    def test_scenario_normal_balance_tuning(self, node):
        """Scenario: typical balance tuning process."""
        # Starting state
        assert node.tune_state == TuneState.IDLE

        # After settle
        node.tune_state = TuneState.TUNING
        assert node.tune_state == TuneState.TUNING

        # After relay measurement
        node.tune_state = TuneState.MEASURING
        node.peaks = [
            (0.0, -0.3),
            (1.5, 0.3),
            (3.0, -0.3),
            (4.5, 0.3),
            (6.0, -0.3),
        ]

        ku, tu = node._analyze_oscillation()
        assert ku is not None
        assert tu is not None

        # Compute gains
        kp, kd = node._compute_ziegler_nichols(ku, tu)
        assert kp > 0
        assert kd > 0

        # Complete
        node.tune_state = TuneState.COMPLETE
        assert node.tune_state == TuneState.COMPLETE
