"""
pose_fusion_node.py — Multi-sensor pose fusion node (Issue #595).

Fuses three sensor streams into a single authoritative map-frame pose:

  IMU (200 Hz)               /imu/data
  UWB + IMU fused pose       /saltybot/pose/fused_cov   (10 Hz)
  Visual odometry            /saltybot/visual_odom       (30 Hz)

The UWB+IMU stream (from saltybot_uwb_imu_fusion) provides absolute position
anchoring every ~100 ms.  Visual odometry fills in smooth relative motion
between UWB updates at 30 Hz.  Raw IMU drives the predict step at 200 Hz.

Publishes
─────────
  /saltybot/pose/authoritative  geometry_msgs/PoseWithCovarianceStamped
  /saltybot/pose/status         std_msgs/String  (JSON, 10 Hz)

TF2
───
  map → base_link  (broadcast at IMU rate when filter is healthy)

Parameters (see config/pose_fusion_params.yaml)
───────────────────────────────────────────────
  Sensor weights:
    sigma_uwb_pos_m      float  0.10   UWB position 1-σ (m)
    use_uwb_covariance   bool   true   extract σ from PoseWithCovarianceStamped
    sigma_uwb_head_rad   float  0.15   UWB heading 1-σ (rad)
    sigma_vo_vel_m_s     float  0.05   VO linear-velocity 1-σ (m/s)
    sigma_vo_omega_r_s   float  0.03   VO yaw-rate 1-σ (rad/s)
    sigma_imu_accel      float  0.05   IMU accel noise (m/s²)
    sigma_imu_gyro       float  0.003  IMU gyro noise (rad/s)
    sigma_vel_drift      float  0.10   velocity process-noise (m/s)
    dropout_vel_damp     float  0.95   velocity damping/s during UWB dropout

  Dropout thresholds:
    uwb_dropout_warn_s   float  2.0    log warning after this many s without UWB
    uwb_dropout_max_s    float  10.0   suppress output after this many s
    vo_dropout_warn_s    float  1.0    log warning after this many s without VO

  Frame IDs:
    map_frame            str    map
    base_frame           str    base_link

  Topics:
    imu_topic            str    /imu/data
    uwb_pose_topic       str    /saltybot/pose/fused_cov
    vo_topic             str    /saltybot/visual_odom
"""

from __future__ import annotations

import json
import math

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

import tf2_ros
from geometry_msgs.msg import (
    PoseWithCovarianceStamped, TransformStamped
)
from nav_msgs.msg import Odometry
from sensor_msgs.msg import Imu
from std_msgs.msg import String

from .pose_fusion_ekf import PoseFusionEKF


# ── QoS ───────────────────────────────────────────────────────────────────────
_SENSOR_QOS = QoSProfile(
    reliability=ReliabilityPolicy.BEST_EFFORT,
    history=HistoryPolicy.KEEP_LAST,
    depth=5,
)
_RELIABLE_QOS = QoSProfile(
    reliability=ReliabilityPolicy.RELIABLE,
    history=HistoryPolicy.KEEP_LAST,
    depth=5,
)


def _quat_to_yaw(qx: float, qy: float, qz: float, qw: float) -> float:
    """Extract yaw (Z-axis rotation) from quaternion."""
    return math.atan2(2.0 * (qw * qz + qx * qy),
                      1.0 - 2.0 * (qy * qy + qz * qz))


def _yaw_to_qz_qw(yaw: float):
    half = yaw * 0.5
    return math.sin(half), math.cos(half)


def _cov36_pos_sigma(cov36: list) -> float | None:
    """Extract average xy position sigma from 6×6 covariance row-major list."""
    try:
        p_xx = cov36[0]
        p_yy = cov36[7]
        if p_xx > 0.0 and p_yy > 0.0:
            return float(math.sqrt((p_xx + p_yy) / 2.0))
    except (IndexError, ValueError):
        pass
    return None


def _cov36_heading_sigma(cov36: list) -> float | None:
    """Extract heading sigma from 6×6 covariance row-major list (index 35 = yaw)."""
    try:
        p_yaw = cov36[35]
        if p_yaw > 0.0:
            return float(math.sqrt(p_yaw))
    except (IndexError, ValueError):
        pass
    return None


class PoseFusionNode(Node):
    """Multi-sensor EKF pose fusion — see module docstring for details."""

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

        # ── Parameters ────────────────────────────────────────────────────────
        self.declare_parameter('sigma_uwb_pos_m',    0.10)
        self.declare_parameter('use_uwb_covariance', True)
        self.declare_parameter('sigma_uwb_head_rad', 0.15)
        self.declare_parameter('sigma_vo_vel_m_s',   0.05)
        self.declare_parameter('sigma_vo_omega_r_s', 0.03)
        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_damp',   0.95)
        self.declare_parameter('uwb_dropout_warn_s', 2.0)
        self.declare_parameter('uwb_dropout_max_s',  10.0)
        self.declare_parameter('vo_dropout_warn_s',  1.0)
        self.declare_parameter('map_frame',           'map')
        self.declare_parameter('base_frame',          'base_link')
        self.declare_parameter('imu_topic',           '/imu/data')
        self.declare_parameter('uwb_pose_topic',      '/saltybot/pose/fused_cov')
        self.declare_parameter('vo_topic',            '/saltybot/visual_odom')
        self.declare_parameter('wheel_odom_topic',    '/saltybot/wheel_odom')
        self.declare_parameter('sigma_wheel_vel_m_s', 0.03)   # wheel vel 1-σ (m/s)
        self.declare_parameter('sigma_wheel_omega_r_s', 0.02) # wheel ωz  1-σ (rad/s)
        self.declare_parameter('use_wheel_odom',      True)

        self._map_frame      = self.get_parameter('map_frame').value
        self._base_frame     = self.get_parameter('base_frame').value
        self._use_uwb_cov    = self.get_parameter('use_uwb_covariance').value
        self._sigma_uwb_pos  = self.get_parameter('sigma_uwb_pos_m').value
        self._sigma_uwb_head = self.get_parameter('sigma_uwb_head_rad').value
        self._uwb_warn       = self.get_parameter('uwb_dropout_warn_s').value
        self._uwb_max        = self.get_parameter('uwb_dropout_max_s').value
        self._vo_warn        = self.get_parameter('vo_dropout_warn_s').value
        self._sigma_wheel_vel   = self.get_parameter('sigma_wheel_vel_m_s').value
        self._sigma_wheel_omega = self.get_parameter('sigma_wheel_omega_r_s').value
        self._use_wheel_odom    = self.get_parameter('use_wheel_odom').value

        # ── EKF ───────────────────────────────────────────────────────────────
        self._ekf = PoseFusionEKF(
            sigma_uwb_pos_m    = self._sigma_uwb_pos,
            sigma_uwb_head_rad = self._sigma_uwb_head,
            sigma_vo_vel_m_s   = self.get_parameter('sigma_vo_vel_m_s').value,
            sigma_vo_omega_r_s = self.get_parameter('sigma_vo_omega_r_s').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_damp   = self.get_parameter('dropout_vel_damp').value,
        )

        self._last_imu_t: float | None = None

        # ── Publishers ─────────────────────────────────────────────────────────
        self._pose_pub   = self.create_publisher(
            PoseWithCovarianceStamped, '/saltybot/pose/authoritative', 10
        )
        self._status_pub = self.create_publisher(
            String, '/saltybot/pose/status', 10
        )
        self._tf_br = tf2_ros.TransformBroadcaster(self)

        # ── Subscriptions ──────────────────────────────────────────────────────
        self.create_subscription(
            Imu,
            self.get_parameter('imu_topic').value,
            self._on_imu,
            _SENSOR_QOS,
        )
        self.create_subscription(
            PoseWithCovarianceStamped,
            self.get_parameter('uwb_pose_topic').value,
            self._on_uwb_pose,
            _RELIABLE_QOS,
        )
        self.create_subscription(
            Odometry,
            self.get_parameter('vo_topic').value,
            self._on_vo,
            _SENSOR_QOS,
        )

        if self._use_wheel_odom:
            self.create_subscription(
                Odometry,
                self.get_parameter('wheel_odom_topic').value,
                self._on_wheel_odom,
                _SENSOR_QOS,
            )

        # ── Status timer (10 Hz) ────────────────────────────────────────────────
        self.create_timer(0.1, self._on_status_timer)

        wheel_status = (
            f'  WheelOdom:{self.get_parameter("wheel_odom_topic").value}'
            if self._use_wheel_odom else ''
        )
        self.get_logger().info(
            f'pose_fusion ready — '
            f'IMU:{self.get_parameter("imu_topic").value}  '
            f'UWB:{self.get_parameter("uwb_pose_topic").value}  '
            f'VO:{self.get_parameter("vo_topic").value}'
            f'{wheel_status}  '
            f'map:{self._map_frame}→{self._base_frame}'
        )

    # ── IMU callback — predict step ────────────────────────────────────────────

    def _on_imu(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

        self._ekf.predict_imu(
            ax_body = msg.linear_acceleration.x,
            ay_body = msg.linear_acceleration.y,
            omega   = msg.angular_velocity.z,
            dt      = dt,
        )

        if not self._ekf.is_initialised:
            return

        # Warn / suppress on UWB dropout
        if self._ekf.uwb_dropout_s > self._uwb_max:
            self.get_logger().warn(
                f'UWB dropout {self._ekf.uwb_dropout_s:.1f}s '
                '> max — pose output suppressed',
                throttle_duration_sec=5.0,
            )
            return

        if self._ekf.uwb_dropout_s > self._uwb_warn:
            self.get_logger().warn(
                f'UWB dropout {self._ekf.uwb_dropout_s:.1f}s — '
                'dead-reckoning only',
                throttle_duration_sec=2.0,
            )

        self._publish_pose(msg.header.stamp)

    # ── UWB+IMU fused pose callback — position+heading update ─────────────────

    def _on_uwb_pose(self, msg: PoseWithCovarianceStamped) -> None:
        x = msg.pose.pose.position.x
        y = msg.pose.pose.position.y

        # Position sigma: from message covariance or parameter
        sigma_pos: float | None = None
        if self._use_uwb_cov:
            sigma_pos = _cov36_pos_sigma(list(msg.pose.covariance))
        if sigma_pos is None or sigma_pos <= 0.0:
            sigma_pos = self._sigma_uwb_pos

        self._ekf.update_uwb_position(x, y, sigma_m=sigma_pos)

        # Heading update from quaternion
        q  = msg.pose.pose.orientation
        yaw = _quat_to_yaw(q.x, q.y, q.z, q.w)

        sigma_head: float | None = None
        if self._use_uwb_cov:
            sigma_head = _cov36_heading_sigma(list(msg.pose.covariance))
        if sigma_head is None or sigma_head <= 0.0:
            sigma_head = self._sigma_uwb_head

        # Only update heading if quaternion is non-trivial (some nodes publish q=(0,0,0,0))
        if abs(q.w) > 0.01 or abs(q.z) > 0.01:
            self._ekf.update_uwb_heading(yaw, sigma_rad=sigma_head)

    # ── Visual-odometry callback — velocity update ─────────────────────────────

    def _on_vo(self, msg: Odometry) -> None:
        if not self._ekf.is_initialised:
            return

        vx_body = msg.twist.twist.linear.x
        omega   = msg.twist.twist.angular.z

        # Extract per-sensor sigma from twist covariance if present
        # Row-major 6×6: [0]=vx, [7]=vy, [35]=ωz
        sigma_vel:   float | None = None
        sigma_omega: float | None = None
        cov = list(msg.twist.covariance)
        if len(cov) == 36:
            p_vx = cov[0]
            p_wz = cov[35]
            if p_vx > 0.0:
                sigma_vel   = float(math.sqrt(p_vx))
            if p_wz > 0.0:
                sigma_omega = float(math.sqrt(p_wz))

        # Warn on VO dropout (but never suppress — VO is secondary)
        if self._ekf.vo_dropout_s > self._vo_warn:
            self.get_logger().warn(
                f'VO dropout {self._ekf.vo_dropout_s:.1f}s',
                throttle_duration_sec=2.0,
            )

        self._ekf.update_vo_velocity(
            vx_body     = vx_body,
            omega       = omega,
            sigma_vel   = sigma_vel,
            sigma_omega = sigma_omega,
        )

    # ── Wheel-odometry callback — velocity update (Issue #646) ────────────────

    def _on_wheel_odom(self, msg: Odometry) -> None:
        """Fuse wheel odometry velocity into EKF (same path as visual odometry)."""
        if not self._ekf.is_initialised:
            return

        vx_body = msg.twist.twist.linear.x
        omega   = msg.twist.twist.angular.z

        # Extract per-axis sigma from twist covariance when available
        sigma_vel:   float | None = None
        sigma_omega: float | None = None
        cov = list(msg.twist.covariance)
        if len(cov) == 36:
            p_vx = cov[0]
            p_wz = cov[35]
            if p_vx > 0.0:
                sigma_vel   = float(math.sqrt(p_vx))
            if p_wz > 0.0:
                sigma_omega = float(math.sqrt(p_wz))

        # Fall back to parameter-configured sigmas
        if sigma_vel is None:
            sigma_vel = self._sigma_wheel_vel
        if sigma_omega is None:
            sigma_omega = self._sigma_wheel_omega

        self._ekf.update_vo_velocity(
            vx_body     = vx_body,
            omega       = omega,
            sigma_vel   = sigma_vel,
            sigma_omega = sigma_omega,
        )

    # ── Status timer (10 Hz) ──────────────────────────────────────────────────

    def _on_status_timer(self) -> None:
        if not self._ekf.is_initialised:
            return
        x, y = self._ekf.position
        status = {
            'x':               round(x, 4),
            'y':               round(y, 4),
            'heading_deg':     round(math.degrees(self._ekf.heading), 2),
            'pos_uncertainty_m': round(self._ekf.position_uncertainty_m(), 4),
            'uwb_dropout_s':   round(self._ekf.uwb_dropout_s, 2),
            'vo_dropout_s':    round(self._ekf.vo_dropout_s,  2),
            'healthy':         self._ekf.uwb_dropout_s < self._uwb_max,
        }
        msg      = String()
        msg.data = json.dumps(status)
        self._status_pub.publish(msg)

    # ── Publish authoritative pose ────────────────────────────────────────────

    def _publish_pose(self, stamp) -> None:
        x, y    = self._ekf.position
        heading = self._ekf.heading
        cov_pos = self._ekf.covariance_position
        cov6    = self._ekf.covariance_full
        vx, vy  = self._ekf.velocity
        omega   = self._ekf.yaw_rate

        qz, qw = _yaw_to_qz_qw(heading)

        msg                          = PoseWithCovarianceStamped()
        msg.header.stamp             = stamp
        msg.header.frame_id          = self._map_frame
        msg.pose.pose.position.x     = x
        msg.pose.pose.position.y     = y
        msg.pose.pose.position.z     = 0.0
        msg.pose.pose.orientation.z  = qz
        msg.pose.pose.orientation.w  = qw

        # Build 6×6 row-major covariance
        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(cov6[2, 2])  # yaw
        msg.pose.covariance = cov36

        self._pose_pub.publish(msg)

        # TF2: map → base_link
        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 = PoseFusionNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.try_shutdown()


if __name__ == '__main__':
    main()
