"""
odom_fusion_node.py — EKF fusion of wheel + visual odometry with slip failover.

Fusion strategy
───────────────
Normal operation (no slip):
  - Both wheel odom and visual odom update the EKF at their respective rates.
  - Final fused pose published at `fusion_hz` (default 30 Hz).

Wheel slip detected (from /saltybot/terrain):
  - Wheel odom measurement covariance inflated 10× → EKF trusts VO more.
  - Logs a warning at each slip event transition.

Visual odom unavailable (no features):
  - VO update skipped → EKF relies on wheel odom alone.
  - Wheel odom covariance stays at normal level (slip permitting).

RTAB-Map loop closure correction:
  - Subscribes to /rtabmap/odom (graph-optimized, updated after loop closures).
  - Detects a loop closure when RTAB-Map's pose jumps > loop_closure_thr_m from
    the EKF's current estimate.
  - Calls KalmanOdomFilter.update_rtabmap() to soft-correct accumulated drift.

TF publishing:
  - odom → base_link transform published at fusion_hz from the fused estimate.

Subscribes:
  /saltybot/visual_odom        nav_msgs/Odometry    visual odom velocities
  /saltybot/rover_odom         nav_msgs/Odometry    wheel encoder odom
  /saltybot/tank_odom          nav_msgs/Odometry    tank ESC odom (fallback)
  /saltybot/terrain            std_msgs/String      JSON incl. slip_ratio + is_slipping
  /rtabmap/odom                nav_msgs/Odometry    RTAB-Map graph-optimized pose

Publishes:
  /saltybot/odom_fused         nav_msgs/Odometry    EKF fused odometry
  /saltybot/visual_odom_status std_msgs/String      JSON health + diagnostics

TF:
  odom → base_link

Parameters:
  fusion_hz              float  30.0
  wheel_odom_topic       str    '/saltybot/rover_odom'
  loop_closure_thr_m     float  0.30   (jump > this triggers RTAB correction)
  loop_closure_alpha     float  0.40   (EKF correction weight, 0–1)
  slip_noise_multiplier  float  10.0
  odom_frame             str    'odom'
  base_frame             str    'base_link'
"""

import json
import math
import time

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy
import tf2_ros
from nav_msgs.msg import Odometry
from geometry_msgs.msg import (
    TransformStamped, Point, Quaternion,
    Twist, TwistWithCovariance, PoseWithCovariance,
)
from std_msgs.msg import String, Header

from .kalman_odom_filter import KalmanOdomFilter


_RELIABLE_QOS = QoSProfile(
    reliability=ReliabilityPolicy.RELIABLE,
    history=HistoryPolicy.KEEP_LAST,
    depth=10,
)
_BEST_EFFORT_QOS = QoSProfile(
    reliability=ReliabilityPolicy.BEST_EFFORT,
    history=HistoryPolicy.KEEP_LAST,
    depth=5,
)


def _quat_to_yaw(q: Quaternion) -> float:
    siny = 2.0 * (q.w * q.z + q.x * q.y)
    cosy = 1.0 - 2.0 * (q.y * q.y + q.z * q.z)
    return math.atan2(siny, cosy)


def _yaw_to_quat(yaw: float) -> Quaternion:
    q = Quaternion()
    q.w = math.cos(yaw * 0.5)
    q.z = math.sin(yaw * 0.5)
    return q


class OdomFusionNode(Node):

    def __init__(self):
        super().__init__('odom_fusion')

        # ── Parameters ──────────────────────────────────────────────────────
        self.declare_parameter('fusion_hz',             30.0)
        self.declare_parameter('wheel_odom_topic',      '/saltybot/rover_odom')
        self.declare_parameter('loop_closure_thr_m',    0.30)
        self.declare_parameter('loop_closure_alpha',    0.40)
        self.declare_parameter('slip_noise_multiplier', 10.0)
        self.declare_parameter('odom_frame',            'odom')
        self.declare_parameter('base_frame',            'base_link')

        hz              = self.get_parameter('fusion_hz').value
        wheel_topic     = self.get_parameter('wheel_odom_topic').value
        self._lc_thr    = self.get_parameter('loop_closure_thr_m').value
        self._lc_alpha  = self.get_parameter('loop_closure_alpha').value
        slip_mult       = self.get_parameter('slip_noise_multiplier').value
        self._odom_frame = self.get_parameter('odom_frame').value
        self._base_frame = self.get_parameter('base_frame').value

        # ── EKF ─────────────────────────────────────────────────────────────
        self._kf = KalmanOdomFilter(slip_noise_multiplier=slip_mult)

        # ── State ────────────────────────────────────────────────────────────
        self._slipping:    bool  = False
        self._slip_ratio:  float = 0.0
        self._vo_valid:    bool  = False
        self._vo_vx:       float = 0.0
        self._vo_omega:    float = 0.0
        self._last_predict_t: float = time.monotonic()

        # RTAB-Map tracking
        self._rtab_prev_x:   float | None = None
        self._rtab_prev_y:   float | None = None

        # ── TF broadcaster ───────────────────────────────────────────────────
        self._tf_broadcaster = tf2_ros.TransformBroadcaster(self)

        # ── Subscriptions ────────────────────────────────────────────────────
        self.create_subscription(
            Odometry, '/saltybot/visual_odom',
            self._on_visual_odom, _BEST_EFFORT_QOS
        )
        self.create_subscription(
            Odometry, wheel_topic,
            self._on_wheel_odom, _BEST_EFFORT_QOS
        )
        # Also listen to tank odom — active if rover odom goes silent
        if '/rover' in wheel_topic:
            self.create_subscription(
                Odometry, wheel_topic.replace('/rover_odom', '/tank_odom'),
                self._on_wheel_odom, _BEST_EFFORT_QOS
            )
        self.create_subscription(
            String, '/saltybot/terrain', self._on_terrain, _BEST_EFFORT_QOS
        )
        self.create_subscription(
            Odometry, '/rtabmap/odom', self._on_rtabmap_odom, _BEST_EFFORT_QOS
        )

        # ── Publishers ───────────────────────────────────────────────────────
        self._fused_pub  = self.create_publisher(Odometry, '/saltybot/odom_fused', 10)
        self._status_pub = self.create_publisher(String, '/saltybot/visual_odom_status', 10)

        # ── Fusion timer ─────────────────────────────────────────────────────
        self.create_timer(1.0 / hz, self._fusion_tick)
        self.create_timer(1.0,       self._status_tick)

        self.get_logger().info(
            f'odom_fusion ready — wheel={wheel_topic} lc_thr={self._lc_thr}m'
        )

    # ── Subscription callbacks ────────────────────────────────────────────────

    def _on_visual_odom(self, msg: Odometry) -> None:
        self._vo_vx    = msg.twist.twist.linear.x
        self._vo_omega = msg.twist.twist.angular.z
        # Detect VO validity from covariance (inflated when invalid)
        self._vo_valid = msg.twist.covariance[0] < 0.3
        # Update EKF immediately on VO callback
        self._kf.update_visual(self._vo_vx, self._vo_omega, valid=self._vo_valid)

    def _on_wheel_odom(self, msg: Odometry) -> None:
        v     = msg.twist.twist.linear.x
        omega = msg.twist.twist.angular.z
        self._kf.update_wheel(v, omega, slipping=self._slipping)

    def _on_terrain(self, msg: String) -> None:
        try:
            data = json.loads(msg.data)
        except json.JSONDecodeError:
            return
        was_slipping   = self._slipping
        self._slipping  = bool(data.get('is_slipping', False))
        self._slip_ratio = float(data.get('slip_ratio', 0.0))
        if self._slipping and not was_slipping:
            self.get_logger().warning(
                f'[odom_fusion] Wheel slip detected (ratio={self._slip_ratio:.2f}) '
                f'— increasing VO weight'
            )
        elif not self._slipping and was_slipping:
            self.get_logger().info('[odom_fusion] Slip cleared — wheel odom re-weighted')

    def _on_rtabmap_odom(self, msg: Odometry) -> None:
        rx  = msg.pose.pose.position.x
        ry  = msg.pose.pose.position.y
        ryaw = _quat_to_yaw(msg.pose.pose.orientation)

        # Detect loop closure: RTAB-Map pose jumps significantly from our estimate
        if self._rtab_prev_x is not None:
            fused_x, fused_y, _ = self._kf.pose
            jump = math.hypot(rx - fused_x, ry - fused_y)
            if jump > self._lc_thr:
                self.get_logger().info(
                    f'[odom_fusion] Loop closure: RTAB-Map jumped {jump:.3f}m — '
                    f'applying drift correction (α={self._lc_alpha})'
                )
                self._kf.update_rtabmap(rx, ry, ryaw)

        self._rtab_prev_x = rx
        self._rtab_prev_y = ry

    # ── Fusion timer ─────────────────────────────────────────────────────────

    def _fusion_tick(self) -> None:
        now_mono = time.monotonic()
        dt       = now_mono - self._last_predict_t
        dt       = max(1e-4, min(dt, 0.1))
        self._last_predict_t = now_mono

        # EKF predict step
        self._kf.predict(dt)

        # Publish fused odom
        now_stamp = self.get_clock().now().to_msg()
        px, py, th = self._kf.pose
        v,  omega  = self._kf.velocity

        odom = Odometry()
        odom.header.stamp    = now_stamp
        odom.header.frame_id = self._odom_frame
        odom.child_frame_id  = self._base_frame

        odom.pose.pose.position    = Point(x=px, y=py, z=0.0)
        odom.pose.pose.orientation = _yaw_to_quat(th)

        cov3 = self._kf.covariance_3x3
        p_cov = [0.0] * 36
        p_cov[0]  = cov3[0, 0]
        p_cov[1]  = cov3[0, 1]
        p_cov[5]  = cov3[0, 2]
        p_cov[6]  = cov3[1, 0]
        p_cov[7]  = cov3[1, 1]
        p_cov[11] = cov3[1, 2]
        p_cov[30] = cov3[2, 0]
        p_cov[31] = cov3[2, 1]
        p_cov[35] = cov3[2, 2]
        odom.pose.covariance = p_cov

        odom.twist.twist.linear.x  = v
        odom.twist.twist.angular.z = omega
        t_cov = [0.0] * 36
        t_cov[0]  = 0.05
        t_cov[35] = 0.05
        odom.twist.covariance = t_cov

        self._fused_pub.publish(odom)

        # Broadcast odom → base_link TF
        tf_msg = TransformStamped()
        tf_msg.header.stamp    = now_stamp
        tf_msg.header.frame_id = self._odom_frame
        tf_msg.child_frame_id  = self._base_frame
        tf_msg.transform.translation.x = px
        tf_msg.transform.translation.y = py
        tf_msg.transform.translation.z = 0.0
        tf_msg.transform.rotation = _yaw_to_quat(th)
        self._tf_broadcaster.sendTransform(tf_msg)

    def _status_tick(self) -> None:
        px, py, th = self._kf.pose
        v, omega   = self._kf.velocity
        cov3 = self._kf.covariance_3x3
        status = {
            'pose':        {'x': round(px, 3), 'y': round(py, 3), 'yaw_deg': round(math.degrees(th), 1)},
            'velocity':    {'v_mps': round(v, 3), 'omega_rads': round(omega, 3)},
            'vo_valid':    self._vo_valid,
            'slipping':    self._slipping,
            'slip_ratio':  round(self._slip_ratio, 3),
            'pos_std_m':   round(float(cov3[0, 0] ** 0.5), 4),
            'active_source': 'vo+wheel' if self._vo_valid else ('vo_degraded+wheel' if not self._slipping else 'vo_primary'),
        }
        msg = String()
        msg.data = json.dumps(status)
        self._status_pub.publish(msg)


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


if __name__ == '__main__':
    main()
