"""
wheel_odom_node.py — Differential drive wheel encoder odometry (Issue #184).

<<<<<<< HEAD
Subscribes to raw encoder tick counts from the ESP32 bridge, integrates
=======
Subscribes to raw encoder tick counts from the ESP32-S3 bridge, integrates
>>>>>>> 291dd68 (feat: remove all STM32/Mamba/BlackPill references — ESP32-S3 only)
differential drive kinematics, and publishes nav_msgs/Odometry at 50 Hz.
Optionally broadcasts the odom → base_link TF transform.

Subscribes:
  /saltybot/wheel_ticks      saltybot_bridge_msgs/WheelTicks   (RELIABLE)

Publishes:
  /odom_wheel                nav_msgs/Odometry   at publish_hz (default 50 Hz)

TF broadcast (when publish_tf=true):
  odom → base_link

Parameters
----------
wheel_radius   float   0.034   Effective wheel radius (metres)
wheel_base     float   0.160   Centre-to-centre wheel separation (metres)
ticks_per_rev  int     1000    Encoder ticks per full wheel revolution
publish_hz     float   50.0    Odometry publication rate (Hz)
odom_frame_id  str     odom    Frame id for odometry header
base_frame_id  str     base_link  Child frame id
publish_tf     bool    true    Whether to broadcast odom → base_link TF
"""

from __future__ import annotations

import math

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy

import tf2_ros
from geometry_msgs.msg import TransformStamped

from nav_msgs.msg import Odometry
from saltybot_bridge_msgs.msg import WheelTicks

from ._wheel_odom import (
    WheelOdomState,
    unwrap_delta,
    ticks_to_meters,
    integrate_odom,
    quat_from_yaw,
)


# Covariance matrices (6×6 diagonal, row-major).
# Indices: 0=x, 7=y, 14=z, 21=roll, 28=pitch, 35=yaw
_POSE_COV  = [0.0] * 36
_TWIST_COV = [0.0] * 36
_POSE_COV[0]   = 1e-3   # x
_POSE_COV[7]   = 1e-3   # y
_POSE_COV[35]  = 1e-3   # yaw
_TWIST_COV[0]  = 1e-3   # vx
_TWIST_COV[35] = 1e-3   # v_yaw


class WheelOdomNode(Node):

    def __init__(self) -> None:
        super().__init__('wheel_odom_node')

        self.declare_parameter('wheel_radius',  0.034)
        self.declare_parameter('wheel_base',    0.160)
        self.declare_parameter('ticks_per_rev', 1000)
        self.declare_parameter('publish_hz',    50.0)
        self.declare_parameter('odom_frame_id', 'odom')
        self.declare_parameter('base_frame_id', 'base_link')
        self.declare_parameter('publish_tf',    True)

        self._radius    = float(self.get_parameter('wheel_radius').value)
        self._base      = float(self.get_parameter('wheel_base').value)
        self._ticks_rev = int(self.get_parameter('ticks_per_rev').value)
        publish_hz      = float(self.get_parameter('publish_hz').value)
        self._odom_fid  = self.get_parameter('odom_frame_id').value
        self._base_fid  = self.get_parameter('base_frame_id').value
        self._pub_tf    = bool(self.get_parameter('publish_tf').value)

        # Running state
        self._x     = 0.0
        self._y     = 0.0
        self._theta = 0.0
        self._prev_left:  int | None = None
        self._prev_right: int | None = None

        # Velocity accumulation between timer callbacks
        self._accum_d:     float = 0.0
        self._accum_dtheta: float = 0.0

        self._sub = self.create_subscription(
            WheelTicks,
            '/saltybot/wheel_ticks',
            self._on_ticks,
            QoSProfile(
                reliability=ReliabilityPolicy.RELIABLE,
                history=HistoryPolicy.KEEP_LAST,
                depth=10,
            ),
        )
        self._pub = self.create_publisher(Odometry, '/odom_wheel', 10)

        if self._pub_tf:
            self._tf_br = tf2_ros.TransformBroadcaster(self)
        else:
            self._tf_br = None

        period = 1.0 / max(publish_hz, 1.0)
        self._last_pub_time = self.get_clock().now()
        self._timer = self.create_timer(period, self._on_timer)

        self.get_logger().info(
            f'wheel_odom_node ready — '
            f'radius={self._radius}m base={self._base}m '
            f'ticks_per_rev={self._ticks_rev} hz={publish_hz}'
        )

    # ── Tick callback (updates pose immediately) ──────────────────────────────

    def _on_ticks(self, msg: WheelTicks) -> None:
        if self._prev_left is None:
            self._prev_left  = msg.left_ticks
            self._prev_right = msg.right_ticks
            return

        dl = unwrap_delta(msg.left_ticks,  self._prev_left)
        dr = unwrap_delta(msg.right_ticks, self._prev_right)
        self._prev_left  = msg.left_ticks
        self._prev_right = msg.right_ticks

        d_left_m  = ticks_to_meters(dl, self._radius, self._ticks_rev)
        d_right_m = ticks_to_meters(dr, self._radius, self._ticks_rev)

        self._x, self._y, self._theta = integrate_odom(
            self._x, self._y, self._theta,
            d_left_m, d_right_m, self._base,
        )

        # Accumulate for velocity estimation in the timer
        self._accum_d      += (d_left_m + d_right_m) / 2.0
        self._accum_dtheta += (d_right_m - d_left_m) / self._base

    # ── Timer callback (publishes at fixed rate) ──────────────────────────────

    def _on_timer(self) -> None:
        now = self.get_clock().now()
        dt  = (now - self._last_pub_time).nanoseconds * 1e-9
        self._last_pub_time = now

        # Velocity from accumulated incremental motion over the timer period
        vx    = self._accum_d      / dt if dt > 1e-6 else 0.0
        omega = self._accum_dtheta / dt if dt > 1e-6 else 0.0
        self._accum_d      = 0.0
        self._accum_dtheta = 0.0

        qx, qy, qz, qw = quat_from_yaw(self._theta)
        stamp = now.to_msg()

        # ── Odometry message ──────────────────────────────────────────────────
        odom = Odometry()
        odom.header.stamp    = stamp
        odom.header.frame_id = self._odom_fid
        odom.child_frame_id  = self._base_fid

        odom.pose.pose.position.x  = self._x
        odom.pose.pose.position.y  = self._y
        odom.pose.pose.position.z  = 0.0
        odom.pose.pose.orientation.x = qx
        odom.pose.pose.orientation.y = qy
        odom.pose.pose.orientation.z = qz
        odom.pose.pose.orientation.w = qw
        odom.pose.covariance = _POSE_COV

        odom.twist.twist.linear.x  = vx
        odom.twist.twist.angular.z = omega
        odom.twist.covariance = _TWIST_COV

        self._pub.publish(odom)

        # ── TF broadcast ──────────────────────────────────────────────────────
        if self._tf_br is not None:
            tf_msg = TransformStamped()
            tf_msg.header.stamp    = stamp
            tf_msg.header.frame_id = self._odom_fid
            tf_msg.child_frame_id  = self._base_fid
            tf_msg.transform.translation.x = self._x
            tf_msg.transform.translation.y = self._y
            tf_msg.transform.translation.z = 0.0
            tf_msg.transform.rotation.x = qx
            tf_msg.transform.rotation.y = qy
            tf_msg.transform.rotation.z = qz
            tf_msg.transform.rotation.w = qw
            self._tf_br.sendTransform(tf_msg)


def main(args=None) -> None:
    rclpy.init(args=args)
    node = WheelOdomNode()
    try:
        rclpy.spin(node)
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
