"""
dynamic_obs_node.py — ROS2 node: LIDAR moving-object detection + Kalman tracking.

Pipeline:
  1. Subscribe /scan (RPLIDAR LaserScan, ~5.5 Hz).
  2. ObjectDetector performs background subtraction → moving blobs.
  3. TrackerManager runs Hungarian assignment + Kalman predict/update at 10 Hz.
  4. Publish /saltybot/moving_objects (MovingObjectArray).
  5. Publish /saltybot/moving_objects_viz (MarkerArray) for RViz.

The 10 Hz timer drives the tracker regardless of scan rate, so prediction
continues between scans (pure-predict steps).

Subscribes:
  /scan                    sensor_msgs/LaserScan   RPLIDAR A1M8

Publishes:
  /saltybot/moving_objects          saltybot_dynamic_obs_msgs/MovingObjectArray
  /saltybot/moving_objects_viz      visualization_msgs/MarkerArray

Parameters:
  max_tracks          int    20
  confirm_frames      int    3
  max_missed_frames   int    6
  assoc_dist_m        float  1.5
  prediction_hz       float  10.0     (tracker + publish rate)
  horizon_s           float  2.5
  pred_step_s         float  0.5
  odom_frame          str    'odom'
  min_speed_mps       float  0.05     (suppress near-zero velocity tracks)
  max_range_m         float  8.0
"""

import time
import math

import numpy as np
import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy

from sensor_msgs.msg import LaserScan
from geometry_msgs.msg import Pose, Point, Quaternion, Vector3
from std_msgs.msg import Header, ColorRGBA
from visualization_msgs.msg import Marker, MarkerArray

try:
    from saltybot_dynamic_obs_msgs.msg import TrackedObject, MovingObjectArray
    _MSGS_AVAILABLE = True
except ImportError:
    _MSGS_AVAILABLE = False

from .object_detector import ObjectDetector
from .tracker_manager import TrackerManager, Track


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


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


class DynamicObsNode(Node):

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

        # ── Parameters ──────────────────────────────────────────────────────
        self.declare_parameter('max_tracks',         20)
        self.declare_parameter('confirm_frames',      3)
        self.declare_parameter('max_missed_frames',   6)
        self.declare_parameter('assoc_dist_m',        1.5)
        self.declare_parameter('prediction_hz',      10.0)
        self.declare_parameter('horizon_s',           2.5)
        self.declare_parameter('pred_step_s',         0.5)
        self.declare_parameter('odom_frame',         'odom')
        self.declare_parameter('min_speed_mps',       0.05)
        self.declare_parameter('max_range_m',         8.0)

        max_tracks    = self.get_parameter('max_tracks').value
        confirm_f     = self.get_parameter('confirm_frames').value
        max_missed    = self.get_parameter('max_missed_frames').value
        assoc_dist    = self.get_parameter('assoc_dist_m').value
        pred_hz       = self.get_parameter('prediction_hz').value
        horizon_s     = self.get_parameter('horizon_s').value
        pred_step     = self.get_parameter('pred_step_s').value
        self._frame   = self.get_parameter('odom_frame').value
        self._min_spd = self.get_parameter('min_speed_mps').value
        self._max_rng = self.get_parameter('max_range_m').value

        # ── Core modules ────────────────────────────────────────────────────
        self._detector = ObjectDetector(
            grid_radius_m=min(self._max_rng + 2.0, 12.0),
            max_cluster=int((self._max_rng / 0.1) ** 2 * 0.5),
        )
        self._tracker = TrackerManager(
            max_tracks=max_tracks,
            confirm_frames=confirm_f,
            max_missed=max_missed,
            assoc_dist_m=assoc_dist,
            horizon_s=horizon_s,
            pred_step_s=pred_step,
        )
        self._horizon_s  = horizon_s
        self._pred_step  = pred_step

        # ── State ────────────────────────────────────────────────────────────
        self._latest_scan: LaserScan | None = None
        self._last_track_t: float = time.monotonic()
        self._scan_processed_stamp: float | None = None

        # ── Subscriptions ────────────────────────────────────────────────────
        self.create_subscription(LaserScan, '/scan', self._on_scan, _SENSOR_QOS)

        # ── Publishers ───────────────────────────────────────────────────────
        if _MSGS_AVAILABLE:
            self._obj_pub = self.create_publisher(
                MovingObjectArray, '/saltybot/moving_objects', 10
            )
        else:
            self._obj_pub = None
            self.get_logger().warning(
                '[dyn_obs] saltybot_dynamic_obs_msgs not built — '
                'MovingObjectArray will not be published'
            )

        self._viz_pub = self.create_publisher(
            MarkerArray, '/saltybot/moving_objects_viz', 10
        )

        # ── Timer ────────────────────────────────────────────────────────────
        self.create_timer(1.0 / pred_hz, self._track_tick)

        self.get_logger().info(
            f'dynamic_obs_tracker ready — '
            f'max_tracks={max_tracks} horizon={horizon_s}s assoc={assoc_dist}m'
        )

    # ── Scan callback ─────────────────────────────────────────────────────────

    def _on_scan(self, msg: LaserScan) -> None:
        self._latest_scan = msg

    # ── 10 Hz tracker tick ────────────────────────────────────────────────────

    def _track_tick(self) -> None:
        t0 = time.monotonic()

        now_mono = t0
        dt = now_mono - self._last_track_t
        dt = max(1e-3, min(dt, 0.5))
        self._last_track_t = now_mono

        scan = self._latest_scan
        detections = []

        if scan is not None:
            stamp_sec = scan.header.stamp.sec + scan.header.stamp.nanosec * 1e-9
            if stamp_sec != self._scan_processed_stamp:
                self._scan_processed_stamp = stamp_sec
                ranges = np.asarray(scan.ranges, dtype=np.float32)
                detections = self._detector.update(
                    ranges,
                    scan.angle_min,
                    scan.angle_increment,
                    min(scan.range_max, self._max_rng),
                )

        confirmed = self._tracker.update(detections, dt)

        latency_ms = (time.monotonic() - t0) * 1000.0
        stamp = self.get_clock().now().to_msg()

        if _MSGS_AVAILABLE and self._obj_pub is not None:
            self._publish_objects(confirmed, stamp, latency_ms)
        self._publish_viz(confirmed, stamp)

    # ── Publish helpers ───────────────────────────────────────────────────────

    def _publish_objects(self, confirmed: list, stamp, latency_ms: float) -> None:
        arr = MovingObjectArray()
        arr.header.stamp    = stamp
        arr.header.frame_id = self._frame
        arr.active_count    = len(confirmed)
        arr.tentative_count = self._tracker.tentative_count
        arr.detector_latency_ms = float(latency_ms)

        for tr in confirmed:
            px, py   = tr.kalman.position
            vx, vy   = tr.kalman.velocity
            speed    = tr.kalman.speed
            if speed < self._min_spd:
                continue

            obj = TrackedObject()
            obj.header          = arr.header
            obj.object_id       = tr.track_id
            obj.pose.pose.position.x = px
            obj.pose.pose.position.y = py
            obj.pose.pose.orientation = _yaw_quat(math.atan2(vy, vx))
            cov = tr.kalman.covariance_2x2
            obj.pose.covariance[0]  = float(cov[0, 0])
            obj.pose.covariance[1]  = float(cov[0, 1])
            obj.pose.covariance[6]  = float(cov[1, 0])
            obj.pose.covariance[7]  = float(cov[1, 1])
            obj.velocity.x  = vx
            obj.velocity.y  = vy
            obj.speed_mps   = speed
            obj.confidence  = min(1.0, tr.hits / (self._tracker._confirm_frames * 3))
            obj.age_frames  = tr.age
            obj.hits        = tr.hits
            obj.is_valid    = True

            # Predicted path
            for px_f, py_f, t_f in tr.kalman.predict_horizon(
                self._horizon_s, self._pred_step
            ):
                p = Pose()
                p.position.x = px_f
                p.position.y = py_f
                p.orientation.w = 1.0
                obj.predicted_path.append(p)
                obj.predicted_times.append(float(t_f))

            arr.objects.append(obj)

        self._obj_pub.publish(arr)

    def _publish_viz(self, confirmed: list, stamp) -> None:
        markers = MarkerArray()

        # Delete old markers
        del_marker = Marker()
        del_marker.header.stamp    = stamp
        del_marker.header.frame_id = self._frame
        del_marker.action = Marker.DELETEALL
        markers.markers.append(del_marker)

        for tr in confirmed:
            px, py = tr.kalman.position
            vx, vy = tr.kalman.velocity
            speed  = tr.kalman.speed
            if speed < self._min_spd:
                continue

            # Cylinder at current position
            m = Marker()
            m.header.stamp    = stamp
            m.header.frame_id = self._frame
            m.ns     = 'dyn_obs'
            m.id     = tr.track_id
            m.type   = Marker.CYLINDER
            m.action = Marker.ADD
            m.pose.position.x = px
            m.pose.position.y = py
            m.pose.position.z = 0.5
            m.pose.orientation.w = 1.0
            m.scale.x = 0.4
            m.scale.y = 0.4
            m.scale.z = 1.0
            m.color   = ColorRGBA(r=1.0, g=0.2, b=0.0, a=0.7)
            m.lifetime.sec = 1
            markers.markers.append(m)

            # Arrow for velocity
            vel_m = Marker()
            vel_m.header  = m.header
            vel_m.ns      = 'dyn_obs_vel'
            vel_m.id      = tr.track_id
            vel_m.type    = Marker.ARROW
            vel_m.action  = Marker.ADD
            from geometry_msgs.msg import Point as GPoint
            p_start = GPoint(x=px, y=py, z=1.0)
            p_end   = GPoint(x=px + vx, y=py + vy, z=1.0)
            vel_m.points  = [p_start, p_end]
            vel_m.scale.x = 0.05
            vel_m.scale.y = 0.10
            vel_m.color   = ColorRGBA(r=1.0, g=1.0, b=0.0, a=0.9)
            vel_m.lifetime.sec = 1
            markers.markers.append(vel_m)

            # Line strip for predicted path
            path_m = Marker()
            path_m.header  = m.header
            path_m.ns      = 'dyn_obs_path'
            path_m.id      = tr.track_id
            path_m.type    = Marker.LINE_STRIP
            path_m.action  = Marker.ADD
            path_m.scale.x = 0.04
            path_m.color   = ColorRGBA(r=1.0, g=0.5, b=0.0, a=0.5)
            path_m.lifetime.sec = 1
            path_m.points.append(GPoint(x=px, y=py, z=0.5))
            for fx, fy, _ in tr.kalman.predict_horizon(self._horizon_s, self._pred_step):
                path_m.points.append(GPoint(x=fx, y=fy, z=0.5))
            markers.markers.append(path_m)

        self._viz_pub.publish(markers)


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


if __name__ == '__main__':
    main()
