"""
ekf_node.py — ROS2 node: UWB + IMU EKF fusion (Issue #573)

Fuses:
  /imu/data              (sensor_msgs/Imu,       200 Hz) — predict step
  /uwb/bearing           (UwbBearing,             10 Hz) — update step (bearing+range → x,y)
  /saltybot/uwb/pose     (geometry_msgs/PoseStamped, 10 Hz) — update step (absolute position)

Publishes:
  /saltybot/pose/fused   (geometry_msgs/PoseStamped)        — fused pose at IMU rate
  /saltybot/pose/fused_cov (geometry_msgs/PoseWithCovarianceStamped) — with covariance

TF2:
  Broadcasts  base_link → map  from fused pose

UWB dropout:
  IMU dead-reckoning continues for up to `max_dead_reckoning_s` seconds.
  After that the velocity estimate is zeroed; position uncertainty grows.

Parameters (see config/fusion_params.yaml):
  imu_topic            /imu/data
  uwb_bearing_topic    /uwb/bearing
  uwb_pose_topic       /saltybot/uwb/pose   (preferred if available)
  use_uwb_pose         true  — use absolute pose; false = use bearing only
  publish_tf           true
  map_frame            map
  base_frame           base_link
  sigma_uwb_pos_m      0.12
  sigma_imu_accel      0.05
  sigma_imu_gyro       0.003
  sigma_vel_drift      0.10
  max_dead_reckoning_s 10.0
"""

import math
import time

import rclpy
from rclpy.node import Node
from rclpy.qos import (
    QoSProfile, ReliabilityPolicy, HistoryPolicy, DurabilityPolicy
)

import tf2_ros
from geometry_msgs.msg import (
    PoseStamped, PoseWithCovarianceStamped, TransformStamped
)
from sensor_msgs.msg import Imu
from std_msgs.msg import Header

from saltybot_uwb_msgs.msg import UwbBearing
from saltybot_uwb_imu_fusion.ekf_math import UwbImuEKF


_SENSOR_QOS = QoSProfile(
    reliability=ReliabilityPolicy.BEST_EFFORT,
    history=HistoryPolicy.KEEP_LAST,
    depth=5,
)


class EkfFusionNode(Node):

    def __init__(self) -> None:
        super().__init__("uwb_imu_fusion")

        # ── Parameters ────────────────────────────────────────────────────────
        self.declare_parameter("imu_topic",             "/imu/data")
        self.declare_parameter("uwb_bearing_topic",     "/uwb/bearing")
        self.declare_parameter("uwb_pose_topic",        "/saltybot/uwb/pose")
        self.declare_parameter("use_uwb_pose",          True)
        self.declare_parameter("publish_tf",            True)
        self.declare_parameter("map_frame",             "map")
        self.declare_parameter("base_frame",            "base_link")
        self.declare_parameter("sigma_uwb_pos_m",       0.12)
        self.declare_parameter("sigma_imu_accel",       0.05)
        self.declare_parameter("sigma_imu_gyro",        0.003)
        self.declare_parameter("sigma_vel_drift",       0.10)
        self.declare_parameter("dropout_vel_damping",   0.95)
        self.declare_parameter("max_dead_reckoning_s",  10.0)

        self._map_frame     = self.get_parameter("map_frame").value
        self._base_frame    = self.get_parameter("base_frame").value
        self._publish_tf    = self.get_parameter("publish_tf").value
        self._use_uwb_pose  = self.get_parameter("use_uwb_pose").value
        self._max_dr        = self.get_parameter("max_dead_reckoning_s").value

        # ── EKF ───────────────────────────────────────────────────────────────
        self._ekf = UwbImuEKF(
            sigma_uwb_pos_m     = self.get_parameter("sigma_uwb_pos_m").value,
            sigma_imu_accel     = self.get_parameter("sigma_imu_accel").value,
            sigma_imu_gyro      = self.get_parameter("sigma_imu_gyro").value,
            sigma_vel_drift     = self.get_parameter("sigma_vel_drift").value,
            dropout_vel_damping = self.get_parameter("dropout_vel_damping").value,
        )

        self._last_imu_t: float | None = None

        # ── Publishers ────────────────────────────────────────────────────────
        self._pose_pub     = self.create_publisher(PoseStamped,                "/saltybot/pose/fused",     10)
        self._pose_cov_pub = self.create_publisher(PoseWithCovarianceStamped,  "/saltybot/pose/fused_cov", 10)

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

        # ── Subscriptions ─────────────────────────────────────────────────────
        self.create_subscription(
            Imu, self.get_parameter("imu_topic").value,
            self._imu_cb, _SENSOR_QOS
        )

        if self._use_uwb_pose:
            self.create_subscription(
                PoseStamped,
                self.get_parameter("uwb_pose_topic").value,
                self._uwb_pose_cb, 10
            )
        else:
            self.create_subscription(
                UwbBearing,
                self.get_parameter("uwb_bearing_topic").value,
                self._uwb_bearing_cb, 10
            )

        self.get_logger().info(
            f"EKF fusion ready — "
            f"IMU:{self.get_parameter('imu_topic').value}  "
            f"UWB:{'pose' if self._use_uwb_pose else 'bearing'}  "
            f"tf={self._publish_tf}  "
            f"dr={self._max_dr}s"
        )

    # ── IMU callback (200 Hz) — predict step ──────────────────────────────────

    def _imu_cb(self, msg: Imu) -> None:
        now = self.get_clock().now().nanoseconds * 1e-9
        if self._last_imu_t is None:
            self._last_imu_t = now
            return

        dt = now - self._last_imu_t
        self._last_imu_t = now

        if dt <= 0.0 or dt > 0.5:
            return  # stale or implausible

        # Extract IMU data (body frame, ROS convention: x=forward, y=left, z=up)
        ax = msg.linear_acceleration.x
        ay = msg.linear_acceleration.y
        # az = msg.linear_acceleration.z  # gravity along z, ignored for ground robot
        omega = msg.angular_velocity.z   # yaw rate

        self._ekf.predict(ax_body=ax, ay_body=ay, omega=omega, dt=dt)

        # Publish fused pose at IMU rate if filter is alive
        if self._ekf.is_initialised:
            if self._ekf.uwb_dropout_s < self._max_dr:
                self._publish_pose(msg.header.stamp)
            else:
                # Dropout too long — stop publishing, log warning
                self.get_logger().warn(
                    f"UWB dropout {self._ekf.uwb_dropout_s:.1f}s — "
                    "pose unreliable, output suppressed",
                    throttle_duration_sec=5.0,
                )

    # ── UWB pose callback (absolute position, 10 Hz) — update step ───────────

    def _uwb_pose_cb(self, msg: PoseStamped) -> None:
        x = msg.pose.position.x
        y = msg.pose.position.y
        self._ekf.update_uwb(x, y)

        if not self._ekf.is_initialised:
            return

        # Extract heading from quaternion if available
        q = msg.pose.orientation
        if abs(q.w) > 0.01:
            yaw = 2.0 * math.atan2(q.z, q.w)
            self._ekf.update_heading(yaw, sigma_rad=0.2)

    # ── UWB bearing callback (relative, 10 Hz) — update step ─────────────────

    def _uwb_bearing_cb(self, msg: UwbBearing) -> None:
        r   = float(msg.range_m)
        brg = float(msg.bearing_rad)
        if r < 0.1:
            return
        # Convert polar (bearing from base_link) to absolute position estimate
        # This is relative to robot, so add to current robot position to get world coords.
        # Note: if we don't have absolute anchors, this stays in base_link frame.
        # Use it as a position measurement in base_link (x = forward, y = left).
        x_rel = r * math.cos(brg)
        y_rel = r * math.sin(brg)
        # Inflate sigma for single-anchor degraded mode
        sigma = 0.15 if msg.confidence >= 1.0 else 0.30
        self._ekf.update_uwb(x_rel, y_rel, sigma_m=sigma)

    # ── Publish fused pose ────────────────────────────────────────────────────

    def _publish_pose(self, stamp) -> None:
        x, y       = self._ekf.position
        heading    = self._ekf.heading
        cov_pos    = self._ekf.covariance_position
        cov_full   = self._ekf.covariance_full

        half_yaw = heading / 2.0
        qz = math.sin(half_yaw)
        qw = math.cos(half_yaw)

        # PoseStamped
        pose = PoseStamped()
        pose.header.stamp    = stamp
        pose.header.frame_id = self._map_frame
        pose.pose.position.x = x
        pose.pose.position.y = y
        pose.pose.position.z = 0.0
        pose.pose.orientation.z = qz
        pose.pose.orientation.w = qw
        self._pose_pub.publish(pose)

        # PoseWithCovarianceStamped
        pose_cov = PoseWithCovarianceStamped()
        pose_cov.header = pose.header
        pose_cov.pose.pose = pose.pose
        cov36 = [0.0] * 36
        cov36[0]  = cov_pos[0, 0]   # x,x
        cov36[1]  = cov_pos[0, 1]   # x,y
        cov36[6]  = cov_pos[1, 0]   # y,x
        cov36[7]  = cov_pos[1, 1]   # y,y
        cov36[14] = 1e-4             # z (not estimated)
        cov36[21] = 1e-4             # roll
        cov36[28] = 1e-4             # pitch
        cov36[35] = float(cov_full[2, 2])  # yaw
        pose_cov.pose.covariance = cov36
        self._pose_cov_pub.publish(pose_cov)

        # TF2
        if self._tf_br is not None:
            tf = TransformStamped()
            tf.header.stamp    = stamp
            tf.header.frame_id = self._map_frame
            tf.child_frame_id  = self._base_frame
            tf.transform.translation.x = x
            tf.transform.translation.y = y
            tf.transform.translation.z = 0.0
            tf.transform.rotation.z    = qz
            tf.transform.rotation.w    = qw
            self._tf_br.sendTransform(tf)


def main(args=None) -> None:
    rclpy.init(args=args)
    node = EkfFusionNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.try_shutdown()


if __name__ == "__main__":
    main()
