"""Unit tests for odom_fusion_node."""

import pytest
import math
import json
import numpy as np
from geometry_msgs.msg import Pose, Quaternion, Twist, Vector3
from nav_msgs.msg import Odometry
from sensor_msgs.msg import Imu
from std_msgs.msg import String, Header
from tf_transformations import quaternion_from_euler, euler_from_quaternion

import rclpy

# Import the node and classes under test
from saltybot_odom_fusion.odom_fusion_node import (
    OdomFusionNode,
    OdomState,
)


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


@pytest.fixture
def node(rclpy_fixture):
    """Create an odometry fusion node instance."""
    node = OdomFusionNode()
    yield node
    node.destroy_node()


class TestOdomState:
    """Test suite for OdomState dataclass."""

    def test_odom_state_initialization(self):
        """Test OdomState initializes with correct defaults."""
        state = OdomState()

        assert state.x == 0.0
        assert state.y == 0.0
        assert state.theta == 0.0
        assert state.vx == 0.0
        assert state.vy == 0.0
        assert state.omega == 0.0

    def test_odom_state_covariance_initialization(self):
        """Test OdomState covariance matrices."""
        state = OdomState()

        assert state.cov_pos.shape == (2, 2)
        assert state.cov_theta == 0.01
        assert state.cov_vel.shape == (2, 2)


class TestOdomFusionNode:
    """Test suite for OdomFusionNode."""

    def test_node_initialization(self, node):
        """Test that node initializes with correct defaults."""
        assert node.state.x == 0.0
        assert node.state.y == 0.0
        assert node.state.theta == 0.0
        assert node.last_wheel_odom is None
        assert node.last_visual_odom is None
        assert node.last_imu is None

    def test_weight_normalization(self, node):
        """Test that weights are normalized to sum to 1.0."""
        total = node.wheel_weight + node.visual_weight + node.imu_weight
        assert abs(total - 1.0) < 1e-6

    def test_wheel_odom_subscription(self, node):
        """Test wheel odometry message handling."""
        odom_msg = Odometry()
        odom_msg.twist.twist.linear.x = 1.0
        odom_msg.twist.twist.linear.y = 0.5

        node._on_wheel_odom(odom_msg)
        assert node.last_wheel_odom is not None
        assert node.last_wheel_odom.twist.twist.linear.x == 1.0

    def test_visual_odom_subscription(self, node):
        """Test visual odometry message handling."""
        odom_msg = Odometry()
        odom_msg.pose.pose.position.x = 5.0
        odom_msg.pose.pose.position.y = 3.0

        node._on_visual_odom(odom_msg)
        assert node.last_visual_odom is not None
        assert node.last_visual_odom.pose.pose.position.x == 5.0

    def test_imu_subscription(self, node):
        """Test IMU message handling."""
        imu_msg = Imu()
        imu_msg.angular_velocity.z = 0.1

        node._on_imu(imu_msg)
        assert node.last_imu is not None
        assert node.last_imu.angular_velocity.z == 0.1

    def test_position_integration(self, node):
        """Test position integration from velocity."""
        node.state.vx = 1.0  # 1 m/s
        node.state.vy = 0.0
        node.state.theta = 0.0  # Facing +x direction

        dt = 1.0  # 1 second
        node._fuse_odometry(dt)

        # Should move 1 meter in +x direction
        assert abs(node.state.x - 1.0) < 0.1

    def test_position_integration_rotated(self, node):
        """Test position integration with rotation."""
        node.state.vx = 1.0  # 1 m/s
        node.state.vy = 0.0
        node.state.theta = math.pi / 2  # Facing +y direction

        dt = 1.0
        node._fuse_odometry(dt)

        # Should move 1 meter in +y direction
        assert abs(node.state.y - 1.0) < 0.1

    def test_angular_velocity_update(self, node):
        """Test angular velocity from IMU."""
        imu_msg = Imu()
        imu_msg.angular_velocity.z = 0.2
        imu_msg.orientation.w = 1.0

        node.last_imu = imu_msg
        dt = 0.1

        node._fuse_odometry(dt)

        # Angular velocity should be weighted by imu_weight
        expected_omega = 0.2 * node.imu_weight
        assert abs(node.state.omega - expected_omega) < 0.01

    def test_velocity_blending_wheel_only(self, node):
        """Test velocity blending with only wheel odometry."""
        wheel_msg = Odometry()
        wheel_msg.twist.twist.linear.x = 1.0
        wheel_msg.twist.twist.linear.y = 0.0

        node.last_wheel_odom = wheel_msg
        dt = 0.1

        node._fuse_odometry(dt)

        # Velocity should be weighted by wheel_weight
        expected_vx = 1.0 * node.wheel_weight
        assert abs(node.state.vx - expected_vx) < 0.01

    def test_visual_drift_correction(self, node):
        """Test drift correction from visual odometry."""
        # Current state
        node.state.x = 0.0
        node.state.y = 0.0

        # Visual odometry says we're at (1.0, 1.0)
        visual_msg = Odometry()
        visual_msg.pose.pose.position.x = 1.0
        visual_msg.pose.pose.position.y = 1.0
        visual_msg.pose.pose.orientation.w = 1.0

        node.last_visual_odom = visual_msg
        dt = 0.1

        node._fuse_odometry(dt)

        # Position should have moved toward visual estimate
        assert node.state.x > 0.0
        assert node.state.y > 0.0

    def test_covariance_updates(self, node):
        """Test that covariances are updated based on measurements."""
        wheel_msg = Odometry()
        visual_msg = Odometry()
        visual_msg.pose.pose.orientation.w = 1.0

        node.last_wheel_odom = wheel_msg
        node.last_visual_odom = visual_msg

        measurements = {
            "wheel": wheel_msg,
            "visual": visual_msg,
        }

        node._update_covariances(measurements)

        # Covariance should be reduced with good measurements
        assert node.state.cov_pos[0, 0] < 0.01

    def test_quaternion_to_yaw(self, node):
        """Test quaternion to yaw conversion."""
        # Create 90 degree yaw rotation
        q = quaternion_from_euler(0, 0, math.pi / 2)

        imu_msg = Imu()
        imu_msg.orientation.x = q[0]
        imu_msg.orientation.y = q[1]
        imu_msg.orientation.z = q[2]
        imu_msg.orientation.w = q[3]

        node.last_imu = imu_msg
        dt = 0.1

        node._fuse_odometry(dt)

        # Theta should be blended toward 90 degrees
        assert node.state.theta > 0.0

    def test_velocity_near_zero(self, node):
        """Test handling of near-zero velocities."""
        wheel_msg = Odometry()
        wheel_msg.twist.twist.linear.x = 0.001  # Very small

        node.last_wheel_odom = wheel_msg
        dt = 0.1

        node._fuse_odometry(dt)

        # Position change should be minimal
        assert abs(node.state.x) < 0.0001

    def test_invalid_dt_ignored(self, node):
        """Test that invalid dt values are ignored."""
        initial_x = node.state.x

        # Zero dt
        node._fuse_odometry(0.0)
        assert node.state.x == initial_x

        # Negative dt
        node._fuse_odometry(-1.0)
        assert node.state.x == initial_x

        # Very large dt
        node._fuse_odometry(10.0)
        assert node.state.x == initial_x


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

    def test_scenario_straight_line(self, node):
        """Scenario: robot moving in straight line."""
        wheel_msg = Odometry()
        wheel_msg.twist.twist.linear.x = 1.0  # 1 m/s forward
        wheel_msg.twist.twist.linear.y = 0.0

        node.last_wheel_odom = wheel_msg
        node.state.theta = 0.0  # Facing +x

        # Move for 5 seconds
        for _ in range(5):
            node._fuse_odometry(1.0)

        # Should be approximately 5 meters in +x direction
        assert 4.0 < node.state.x < 6.0
        assert abs(node.state.y) < 0.5

    def test_scenario_circular_motion(self, node):
        """Scenario: robot moving in circle."""
        wheel_msg = Odometry()
        wheel_msg.twist.twist.linear.x = 1.0  # 1 m/s forward

        imu_msg = Imu()
        imu_msg.angular_velocity.z = 0.5  # 0.5 rad/s rotation
        imu_msg.orientation.w = 1.0

        node.last_wheel_odom = wheel_msg
        node.last_imu = imu_msg

        # Simulate circular motion for 4 cycles
        for i in range(4):
            node._fuse_odometry(1.0)

        # Should have moved distance and rotated
        distance = math.sqrt(node.state.x**2 + node.state.y**2)
        assert distance > 1.0

    def test_scenario_drift_correction(self, node):
        """Scenario: wheel odometry drifts, vision corrects."""
        # Wheel says we're at 10, 10
        node.state.x = 10.0
        node.state.y = 10.0

        # Vision says we're actually at 9, 9
        visual_msg = Odometry()
        visual_msg.pose.pose.position.x = 9.0
        visual_msg.pose.pose.position.y = 9.0
        visual_msg.pose.pose.orientation.w = 1.0

        node.last_visual_odom = visual_msg

        # Before correction
        initial_x = node.state.x
        initial_y = node.state.y

        node._fuse_odometry(1.0)

        # Position should move toward visual estimate
        assert node.state.x < initial_x
        assert node.state.y < initial_y

    def test_scenario_all_sensors_available(self, node):
        """Scenario: all three sensor sources available."""
        wheel_msg = Odometry()
        wheel_msg.twist.twist.linear.x = 0.5

        visual_msg = Odometry()
        visual_msg.pose.pose.position.x = 2.0
        visual_msg.pose.pose.position.y = 1.0
        visual_msg.pose.pose.orientation.w = 1.0

        imu_msg = Imu()
        q = quaternion_from_euler(0, 0, 0.1)  # Small yaw
        imu_msg.orientation.x = q[0]
        imu_msg.orientation.y = q[1]
        imu_msg.orientation.z = q[2]
        imu_msg.orientation.w = q[3]
        imu_msg.angular_velocity.z = 0.05

        node.last_wheel_odom = wheel_msg
        node.last_visual_odom = visual_msg
        node.last_imu = imu_msg

        measurements = {
            "wheel": wheel_msg,
            "visual": visual_msg,
            "imu": imu_msg,
        }

        node._update_covariances(measurements)

        # Covariance should be low with all sources
        assert node.state.cov_pos[0, 0] < 0.01
        assert node.state.cov_theta < 0.01

    def test_scenario_sensor_dropout(self, node):
        """Scenario: visual sensor drops out, wheel continues."""
        wheel_msg = Odometry()
        wheel_msg.twist.twist.linear.x = 1.0

        node.last_wheel_odom = wheel_msg
        node.last_visual_odom = None  # Visual dropout
        node.last_imu = None

        # Should continue with wheel only
        node._fuse_odometry(1.0)

        # Position should still integrate
        assert node.state.x > 0.0
