"""
tracking_fusion_node.py — Multi-modal tracking fusion for saltybot.

Subscribes
----------
  /uwb/target     (geometry_msgs/PoseStamped)  — UWB-triangulated position  (~10 Hz)
  /person/target  (geometry_msgs/PoseStamped)  — camera-detected position   (~30 Hz)

Publishes
---------
  /social/tracking/fused_target  (saltybot_social_msgs/FusedTarget)  at control_rate Hz

Algorithm
---------
  1. Each incoming measurement updates a 4-state Kalman filter [x, y, vx, vy].
  2. A 20 Hz timer runs predict+select+publish:
       a. KF predict(dt)
       b. Compute per-source confidence from measurement age + staleness model
       c. If either source is live:
            - "fused"  → confidence-weighted position blend → KF update
            - "uwb"    → UWB position → KF update
            - "camera" → camera position → KF update
       d. Build FusedTarget from KF state + composite confidence
  3. If all sources are lost but within predict_timeout, keep publishing with
     active_source="predicted" and degrading confidence.
  4. Beyond predict_timeout, no message is published (node stays alive).

Kalman tuning for EUC speeds (20–30 km/h ≈ 5.5–8.3 m/s)
---------------------------------------------------------
  process_noise=3.0 m/s²  — allows rapid acceleration events
  predict_timeout=3.0 s    — coasts ≈30 m at 10 m/s; acceptable dead-reckoning

Parameters
----------
  uwb_timeout         : UWB staleness threshold (s)            default 1.5
  cam_timeout         : Camera staleness threshold (s)          default 1.0
  predict_timeout     : Max KF coast before no publish (s)      default 3.0
  process_noise       : KF acceleration noise std-dev (m/s²)    default 3.0
  meas_noise_uwb      : UWB position noise std-dev (m)          default 0.20
  meas_noise_cam      : Camera position noise std-dev (m)       default 0.12
  control_rate        : Publish / KF predict rate (Hz)          default 20.0
  confidence_threshold: Min source confidence to use (0–1)      default 0.15

Usage
-----
  ros2 launch saltybot_social_tracking tracking.launch.py
"""

import math
import time

import rclpy
from rclpy.node import Node
from geometry_msgs.msg import PoseStamped
from std_msgs.msg import Header

from saltybot_social_msgs.msg import FusedTarget
from saltybot_social_tracking.kalman_tracker import KalmanTracker
from saltybot_social_tracking.source_arbiter import (
    uwb_confidence,
    camera_confidence,
    select_source,
    fuse_positions,
    composite_confidence,
    bearing_and_range,
)

_BIG_AGE = 1e9   # sentinel: source never received


class TrackingFusionNode(Node):

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

        # ── Parameters ────────────────────────────────────────────────────────
        self.declare_parameter("uwb_timeout",          1.5)
        self.declare_parameter("cam_timeout",          1.0)
        self.declare_parameter("predict_timeout",      3.0)
        self.declare_parameter("process_noise",        3.0)
        self.declare_parameter("meas_noise_uwb",       0.20)
        self.declare_parameter("meas_noise_cam",       0.12)
        self.declare_parameter("control_rate",        20.0)
        self.declare_parameter("confidence_threshold", 0.15)

        self._p = self._load_params()

        # ── State ─────────────────────────────────────────────────────────────
        self._kf = KalmanTracker(
            process_noise=self._p["process_noise"],
            meas_noise_uwb=self._p["meas_noise_uwb"],
            meas_noise_cam=self._p["meas_noise_cam"],
        )
        self._last_uwb_t: float = 0.0     # monotonic; 0 = never received
        self._last_cam_t: float = 0.0
        self._uwb_x: float = 0.0
        self._uwb_y: float = 0.0
        self._cam_x: float = 0.0
        self._cam_y: float = 0.0
        self._uwb_tag_id: str = ""
        self._last_predict_t: float = 0.0  # monotonic time of last predict call
        self._last_any_t: float = 0.0      # monotonic time of last live measurement

        # ── Subscriptions ─────────────────────────────────────────────────────
        self.create_subscription(
            PoseStamped, "/uwb/target", self._uwb_cb, 10)
        self.create_subscription(
            PoseStamped, "/person/target", self._cam_cb, 10)

        # ── Publisher ─────────────────────────────────────────────────────────
        self._pub = self.create_publisher(FusedTarget, "/social/tracking/fused_target", 10)

        # ── Timer ─────────────────────────────────────────────────────────────
        self._timer = self.create_timer(
            1.0 / self._p["control_rate"], self._control_cb)

        self.get_logger().info(
            f"TrackingFusion ready  "
            f"rate={self._p['control_rate']}Hz  "
            f"uwb_timeout={self._p['uwb_timeout']}s  "
            f"cam_timeout={self._p['cam_timeout']}s  "
            f"predict_timeout={self._p['predict_timeout']}s  "
            f"process_noise={self._p['process_noise']}m/s²"
        )

    # ── Parameter helpers ──────────────────────────────────────────────────────

    def _load_params(self) -> dict:
        return {
            "uwb_timeout":          self.get_parameter("uwb_timeout").value,
            "cam_timeout":          self.get_parameter("cam_timeout").value,
            "predict_timeout":      self.get_parameter("predict_timeout").value,
            "process_noise":        self.get_parameter("process_noise").value,
            "meas_noise_uwb":       self.get_parameter("meas_noise_uwb").value,
            "meas_noise_cam":       self.get_parameter("meas_noise_cam").value,
            "control_rate":         self.get_parameter("control_rate").value,
            "confidence_threshold": self.get_parameter("confidence_threshold").value,
        }

    # ── Measurement callbacks ──────────────────────────────────────────────────

    def _uwb_cb(self, msg: PoseStamped) -> None:
        self._uwb_x = msg.pose.position.x
        self._uwb_y = msg.pose.position.y
        self._uwb_tag_id = ""  # PoseStamped has no tag field; tag reported via /uwb/bearing
        t = time.monotonic()
        self._last_uwb_t = t
        self._last_any_t = t

        # Seed KF on first measurement
        if not self._kf.initialized:
            self._kf.initialize(self._uwb_x, self._uwb_y)
            self._last_predict_t = t

    def _cam_cb(self, msg: PoseStamped) -> None:
        self._cam_x = msg.pose.position.x
        self._cam_y = msg.pose.position.y
        t = time.monotonic()
        self._last_cam_t = t
        self._last_any_t = t

        if not self._kf.initialized:
            self._kf.initialize(self._cam_x, self._cam_y)
            self._last_predict_t = t

    # ── Control loop ───────────────────────────────────────────────────────────

    def _control_cb(self) -> None:
        self._p = self._load_params()

        if not self._kf.initialized:
            return   # no data yet — nothing to publish

        now = time.monotonic()
        dt = now - self._last_predict_t if self._last_predict_t > 0.0 else (
            1.0 / self._p["control_rate"]
        )
        self._last_predict_t = now

        # KF predict
        self._kf.predict(dt)

        # Source confidence
        uwb_age = (now - self._last_uwb_t) if self._last_uwb_t > 0.0 else _BIG_AGE
        cam_age = (now - self._last_cam_t) if self._last_cam_t > 0.0 else _BIG_AGE

        u_conf = uwb_confidence(uwb_age, self._p["uwb_timeout"])
        c_conf = camera_confidence(cam_age, self._p["cam_timeout"])

        threshold = self._p["confidence_threshold"]
        source = select_source(u_conf, c_conf, threshold)

        if source == "predicted":
            # Check predict_timeout — stop publishing if too stale
            last_live_age = (
                (now - self._last_any_t) if self._last_any_t > 0.0 else _BIG_AGE
            )
            if last_live_age > self._p["predict_timeout"]:
                return   # silently stop publishing

        # Apply measurement update if a live source exists
        if source == "fused":
            fx, fy = fuse_positions(
                self._uwb_x, self._uwb_y, u_conf,
                self._cam_x, self._cam_y, c_conf,
            )
            self._kf.update(fx, fy, source="uwb")   # use tighter noise for blended
        elif source == "uwb":
            self._kf.update(self._uwb_x, self._uwb_y, source="uwb")
        elif source == "camera":
            self._kf.update(self._cam_x, self._cam_y, source="camera")
        # "predicted" → no update; KF coasts

        # Build and publish message
        kx, ky = self._kf.position
        vx, vy = self._kf.velocity
        kf_unc = self._kf.position_uncertainty_m()
        conf = composite_confidence(u_conf, c_conf, source, kf_unc)
        bearing, range_m = bearing_and_range(kx, ky)

        hdr = Header()
        hdr.stamp    = self.get_clock().now().to_msg()
        hdr.frame_id = "base_link"

        msg = FusedTarget()
        msg.header         = hdr
        msg.position.x     = kx
        msg.position.y     = ky
        msg.position.z     = 0.0
        msg.velocity.x     = vx
        msg.velocity.y     = vy
        msg.velocity.z     = 0.0
        msg.range_m        = float(range_m)
        msg.bearing_rad    = float(bearing)
        msg.confidence     = float(conf)
        msg.active_source  = source
        msg.tag_id         = self._uwb_tag_id if "uwb" in source else ""

        self._pub.publish(msg)


# ── Entry point ────────────────────────────────────────────────────────────────

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


if __name__ == "__main__":
    main()
