"""
person_reid_node.py — Person re-identification for cross-camera tracking.

Subscribes to:
  /person/detections      vision_msgs/Detection2DArray  (person bounding boxes)
  /camera/color/image_raw sensor_msgs/Image             (colour frame for crops)

Publishes:
  /saltybot/person_reid   saltybot_person_reid_msgs/PersonAppearanceArray  (5 Hz)

For each detected person the node:
  1. Crops the torso region (top 65 % of the bounding box height).
  2. Extracts a 128-dim L2-normalised embedding via MobileNetV2 ONNX (if the
     model file is provided) or a 128-bin HSV colour histogram (fallback).
  3. Matches against a cosine-similarity gallery.
  4. Assigns a persistent identity_id (new or existing).

Parameters:
  model_path          str    ''      Path to MobileNetV2+projection ONNX file
  match_threshold     float  0.75    Cosine similarity threshold for matching
  max_identity_age_s  float  300.0   Seconds before an unseen identity is pruned
  publish_hz          float  5.0     Publication rate (Hz)
"""

from __future__ import annotations

from typing import List

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

import message_filters
import cv2
import numpy as np
from cv_bridge import CvBridge

from sensor_msgs.msg import Image
from vision_msgs.msg import Detection2DArray

from saltybot_person_reid_msgs.msg import PersonAppearance, PersonAppearanceArray

from ._embedding_model import EmbeddingModel
from ._reid_gallery import ReidGallery

# Fraction of bbox height kept as torso crop (top portion)
_TORSO_FRAC = 0.65

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


class PersonReidNode(Node):

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

        self.declare_parameter('model_path',         '')
        self.declare_parameter('match_threshold',    0.75)
        self.declare_parameter('max_identity_age_s', 300.0)
        self.declare_parameter('publish_hz',         5.0)

        model_path  = self.get_parameter('model_path').value
        match_thr   = self.get_parameter('match_threshold').value
        max_age     = self.get_parameter('max_identity_age_s').value
        publish_hz  = self.get_parameter('publish_hz').value

        self._bridge   = CvBridge()
        self._embedder = EmbeddingModel(model_path or None)
        self._gallery  = ReidGallery(match_threshold=match_thr, max_age_s=max_age)

        # Buffer: updated by frame callback, drained by timer
        self._pending: List[PersonAppearance] = []
        self._pending_header = None

        # Synchronized subscribers
        det_sub = message_filters.Subscriber(
            self, Detection2DArray, '/person/detections',
            qos_profile=_BEST_EFFORT_QOS)
        img_sub = message_filters.Subscriber(
            self, Image, '/camera/color/image_raw',
            qos_profile=_BEST_EFFORT_QOS)
        self._sync = message_filters.ApproximateTimeSynchronizer(
            [det_sub, img_sub], queue_size=4, slop=0.1)
        self._sync.registerCallback(self._on_frame)

        self._pub = self.create_publisher(
            PersonAppearanceArray, '/saltybot/person_reid', 10)

        self.create_timer(1.0 / publish_hz, self._tick)

        backend = 'ONNX' if self._embedder._net else 'histogram'
        self.get_logger().info(
            f'person_reid ready — backend={backend} '
            f'threshold={match_thr} max_age={max_age}s'
        )

    # ── Frame callback ─────────────────────────────────────────────────────────

    def _on_frame(self, det_msg: Detection2DArray, img_msg: Image) -> None:
        if not det_msg.detections:
            self._pending = []
            self._pending_header = det_msg.header
            return

        try:
            bgr = self._bridge.imgmsg_to_cv2(img_msg, desired_encoding='bgr8')
        except Exception as exc:
            self.get_logger().error(
                f'imgmsg_to_cv2 failed: {exc}', throttle_duration_sec=5.0)
            return

        h_img, w_img = bgr.shape[:2]
        appearances: List[PersonAppearance] = []

        for det in det_msg.detections:
            cx  = det.bbox.center.position.x
            cy  = det.bbox.center.position.y
            bw  = det.bbox.size_x
            bh  = det.bbox.size_y
            conf = det.results[0].hypothesis.score if det.results else 0.0

            # Torso crop: top TORSO_FRAC of bounding box
            x1 = max(0,      int(cx - bw / 2.0))
            y1 = max(0,      int(cy - bh / 2.0))
            x2 = min(w_img,  int(cx + bw / 2.0))
            y2 = min(h_img,  int(cy - bh / 2.0 + bh * _TORSO_FRAC))

            if x2 - x1 < 8 or y2 - y1 < 8:
                continue

            crop = bgr[y1:y2, x1:x2]
            emb  = self._embedder.embed(crop)
            identity_id, match_score, is_new = self._gallery.match(emb)

            app = PersonAppearance()
            app.header          = det_msg.header
            app.track_id        = identity_id
            app.embedding       = emb.tolist()
            app.bbox            = det.bbox
            app.confidence      = float(conf)
            app.match_score     = float(match_score)
            app.is_new_identity = is_new
            appearances.append(app)

        self._pending = appearances
        self._pending_header = det_msg.header

    # ── 5 Hz publish tick ─────────────────────────────────────────────────────

    def _tick(self) -> None:
        if self._pending_header is None:
            return
        msg = PersonAppearanceArray()
        msg.header      = self._pending_header
        msg.appearances = self._pending
        self._pub.publish(msg)


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


if __name__ == '__main__':
    main()
