"""
landmark_smoother_node.py — Per-person 68-point Kalman landmark smoother (Issue #227).

Subscribes to raw 68-point facial landmarks and republishes a smoothed version
using a constant-velocity Kalman filter maintained per face track.

Subscribes:
  /social/faces/landmarks       saltybot_social_msgs/FaceLandmarksArray

Publishes:
  /social/faces/landmarks_smooth  saltybot_social_msgs/FaceLandmarksArray

Algorithm
---------
1. For each FaceLandmarks message, look up (or create) a per-track
   LandmarkKalmanFilter keyed by person_id (or face_id when person_id == -1).
2. Call kf.update(obs, dt) → smoothed (68, 2) positions.
3. Publish immediately — no extra timer delay.
4. A 1 Hz timer prunes tracks not updated within `prune_age_s` (default 3.0 s).

Parameters
----------
prune_age_s       float  3.0   Seconds of inactivity before a track is dropped
process_noise_std float  0.01  Kalman process noise std-dev (normalised coords/s)
obs_noise_std     float  0.005 Kalman observation noise std-dev (normalised coords)
"""

from __future__ import annotations

import time
from typing import Dict, Optional, Tuple

import numpy as np

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy

from saltybot_social_msgs.msg import FaceLandmarks, FaceLandmarksArray

from ._landmark_kalman import LandmarkKalmanFilter


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


class LandmarkSmootherNode(Node):

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

        self.declare_parameter('prune_age_s',       3.0)
        self.declare_parameter('process_noise_std', 0.01)
        self.declare_parameter('obs_noise_std',     0.005)

        self._prune_age   = self.get_parameter('prune_age_s').value
        self._proc_noise  = self.get_parameter('process_noise_std').value
        self._obs_noise   = self.get_parameter('obs_noise_std').value

        # track_key → (LandmarkKalmanFilter, last_update_time, last_stamp_sec)
        self._tracks: Dict[int, Tuple[LandmarkKalmanFilter, float, float]] = {}

        self._sub = self.create_subscription(
            FaceLandmarksArray,
            '/social/faces/landmarks',
            self._on_landmarks,
            _SENSOR_QOS,
        )
        self._pub = self.create_publisher(
            FaceLandmarksArray,
            '/social/faces/landmarks_smooth',
            10,
        )
        self.create_timer(1.0, self._prune_stale_tracks)

        self.get_logger().info(
            f'landmark_smoother_node ready — '
            f'prune_age={self._prune_age}s '
            f'q={self._proc_noise} r={self._obs_noise}'
        )

    # ── Callback ──────────────────────────────────────────────────────────────

    def _on_landmarks(self, msg: FaceLandmarksArray) -> None:
        now = time.monotonic()
        out = FaceLandmarksArray()
        out.header = msg.header

        for face in msg.faces:
            # Validate landmark count
            if len(face.landmarks) != 136:
                self.get_logger().warn(
                    f'face_id={face.face_id} has {len(face.landmarks)} floats, '
                    f'expected 136 — skipping',
                    throttle_duration_sec=5.0,
                )
                out.faces.append(face)   # pass through unchanged
                continue

            obs = np.array(face.landmarks, dtype=np.float64).reshape(68, 2)
            track_key = self._track_key(face)

            if track_key in self._tracks:
                kf, _last_t, last_stamp = self._tracks[track_key]
                msg_stamp = _stamp_to_sec(msg.header.stamp)
                dt = msg_stamp - last_stamp if msg_stamp > last_stamp else 1.0 / 30.0
                smoothed = kf.update(obs, dt)
            else:
                kf = LandmarkKalmanFilter(
                    n_points=68,
                    process_noise_std=self._proc_noise,
                    obs_noise_std=self._obs_noise,
                )
                smoothed = kf.update(obs, dt=0.0)

            msg_stamp = _stamp_to_sec(msg.header.stamp)
            self._tracks[track_key] = (kf, now, msg_stamp)

            smoothed_face = FaceLandmarks()
            smoothed_face.header    = face.header
            smoothed_face.face_id   = face.face_id
            smoothed_face.person_id = face.person_id
            smoothed_face.confidence = face.confidence
            smoothed_face.landmarks  = smoothed.flatten().tolist()
            out.faces.append(smoothed_face)

        self._pub.publish(out)

    # ── Helpers ───────────────────────────────────────────────────────────────

    @staticmethod
    def _track_key(face: FaceLandmarks) -> int:
        """Use person_id when valid, otherwise fall back to face_id."""
        return face.person_id if face.person_id >= 0 else -(face.face_id + 1)

    def _prune_stale_tracks(self) -> None:
        now = time.monotonic()
        stale = [k for k, (_, t, _) in self._tracks.items()
                 if now - t > self._prune_age]
        for k in stale:
            del self._tracks[k]
        if stale:
            self.get_logger().debug(f'Pruned {len(stale)} stale track(s)')


def _stamp_to_sec(stamp) -> float:
    return stamp.sec + stamp.nanosec * 1e-9


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


if __name__ == '__main__':
    main()
