#!/usr/bin/env python3
"""Tilt-compensated compass heading node for SaltyBot.

Magnetometer heading with declination correction and tilt compensation.
France magnetic declination: 1.5 degrees east.

Published topics:
  /saltybot/heading (std_msgs/Float64) - Heading in degrees (0-360)
  /saltybot/heading_quaternion (geometry_msgs/QuaternionStamped) - Heading as Z-axis quaternion

Subscribed topics:
  /saltybot/imu/data (sensor_msgs/Imu) - IMU with magnetometer and orientation
  /saltybot/mag (sensor_msgs/MagneticField) - Magnetometer data
"""

import math
from typing import Optional

import rclpy
from rclpy.node import Node
from rclpy.timer import Timer
from geometry_msgs.msg import Quaternion, QuaternionStamped
from std_msgs.msg import Float64
from sensor_msgs.msg import Imu, MagneticField


class CompassHeadingNode(Node):
    """ROS2 node for tilt-compensated compass heading."""

    # France magnetic declination: 1.5 degrees east
    MAGNETIC_DECLINATION_DEG = 1.5

    def __init__(self):
        super().__init__("compass_heading")

        # Parameters
        self.declare_parameter("frequency", 10)  # Hz
        frequency = self.get_parameter("frequency").value
        self.declare_parameter("declination_deg", self.MAGNETIC_DECLINATION_DEG)
        self.declination_rad = math.radians(
            self.get_parameter("declination_deg").value
        )

        # Latest sensor data
        self.mag_field: Optional[MagneticField] = None
        self.imu_data: Optional[Imu] = None

        # Subscriptions
        self.create_subscription(
            Imu, "/saltybot/imu/data", self._on_imu_data, 10
        )
        self.create_subscription(
            MagneticField, "/saltybot/mag", self._on_mag_field, 10
        )

        # Publications
        self.pub_heading = self.create_publisher(Float64, "/saltybot/heading", 10)
        self.pub_heading_quat = self.create_publisher(
            QuaternionStamped, "/saltybot/heading_quaternion", 10
        )

        # Timer for heading computation at specified frequency
        period = 1.0 / frequency
        self.timer: Timer = self.create_timer(period, self._timer_callback)

        self.get_logger().info(
            f"Compass heading node initialized at {frequency}Hz. "
            f"Declination: {self.get_parameter('declination_deg').value}°"
        )

    def _on_imu_data(self, msg: Imu) -> None:
        """Update IMU data (orientation for tilt compensation)."""
        self.imu_data = msg

    def _on_mag_field(self, msg: MagneticField) -> None:
        """Update magnetometer field data."""
        self.mag_field = msg

    def _timer_callback(self) -> None:
        """Compute and publish tilt-compensated heading at 10Hz."""
        if self.mag_field is None or self.imu_data is None:
            return

        # Extract magnetometer components (Tesla)
        mag_x = self.mag_field.magnetic_field.x
        mag_y = self.mag_field.magnetic_field.y
        mag_z = self.mag_field.magnetic_field.z

        # Extract quaternion from IMU (orientation in world frame)
        q = self.imu_data.orientation
        qx, qy, qz, qw = q.x, q.y, q.z, q.w

        # Convert quaternion to roll, pitch, yaw
        roll, pitch, _ = self._quaternion_to_euler(qx, qy, qz, qw)

        # Tilt-compensated heading calculation
        # Project magnetometer vector onto horizontal plane accounting for roll and pitch
        mag_north = (
            mag_x * math.cos(pitch)
            + mag_y * math.sin(roll) * math.sin(pitch)
            + mag_z * math.cos(roll) * math.sin(pitch)
        )
        mag_east = mag_y * math.cos(roll) - mag_z * math.sin(roll)

        # Compute heading from north and east components
        heading_rad = math.atan2(mag_east, mag_north)

        # Apply magnetic declination correction (add declination to convert magnetic north to true north)
        heading_corrected_rad = heading_rad + self.declination_rad

        # Normalize to 0-360 degrees
        heading_deg = math.degrees(heading_corrected_rad)
        if heading_deg < 0:
            heading_deg += 360.0
        if heading_deg >= 360.0:
            heading_deg -= 360.0

        # Publish heading as Float64
        heading_msg = Float64()
        heading_msg.data = heading_deg
        self.pub_heading.publish(heading_msg)

        # Create quaternion for heading-only rotation (around Z axis)
        # Heading angle is around Z axis, roll and pitch are zero for heading-only representation
        heading_quat = self._create_heading_quaternion(heading_rad)

        # Publish heading as quaternion (stamped)
        quat_msg = QuaternionStamped()
        quat_msg.header.stamp = self.get_clock().now().to_msg()
        quat_msg.header.frame_id = "map"
        quat_msg.quaternion = heading_quat

        self.pub_heading_quaternion.publish(quat_msg)

    def _quaternion_to_euler(
        self, qx: float, qy: float, qz: float, qw: float
    ) -> tuple:
        """Convert quaternion to roll, pitch, yaw in radians."""
        # Roll (rotation around x-axis)
        sinr_cosp = 2 * (qw * qx + qy * qz)
        cosr_cosp = 1 - 2 * (qx * qx + qy * qy)
        roll = math.atan2(sinr_cosp, cosr_cosp)

        # Pitch (rotation around y-axis)
        sinp = 2 * (qw * qy - qz * qx)
        if abs(sinp) >= 1:
            pitch = math.copysign(math.pi / 2, sinp)
        else:
            pitch = math.asin(sinp)

        # Yaw (rotation around z-axis)
        siny_cosp = 2 * (qw * qz + qx * qy)
        cosy_cosp = 1 - 2 * (qy * qy + qz * qz)
        yaw = math.atan2(siny_cosp, cosy_cosp)

        return roll, pitch, yaw

    def _create_heading_quaternion(self, heading_rad: float) -> Quaternion:
        """Create a quaternion representing rotation around Z axis (heading only)."""
        # Quaternion for rotation around Z axis: q = [0, 0, sin(θ/2), cos(θ/2)]
        half_angle = heading_rad / 2.0
        quat = Quaternion()
        quat.x = 0.0
        quat.y = 0.0
        quat.z = math.sin(half_angle)
        quat.w = math.cos(half_angle)
        return quat


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


if __name__ == "__main__":
    main()
