#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from rclpy.executors import MultiThreadedExecutor
from rclpy.callback_groups import ReentrantCallbackGroup, MutuallyExclusiveCallbackGroup
from cv_bridge import CvBridge, CvBridgeError
import torch
import torch.nn.functional as F
import torchvision
from torchvision import transforms
import clip
import cv2
import json
import os
import time
from datetime import datetime
import numpy as np
from ament_index_python.packages import get_package_share_directory
from tf2_ros import Buffer, TransformListener
from rclpy.duration import Duration
from geometry_msgs.msg import PoseStamped
from sensor_msgs.msg import Image, CameraInfo
import tf2_geometry_msgs
from tf2_geometry_msgs.tf2_geometry_msgs import do_transform_pose




    

class ClipDetectionNode(Node):
    def __init__(self):
        super().__init__('clip_detector')
        self.declare_parameter('reset_records_on_startup', False)
        self.declare_parameter('detection_interval_sec', 1.0)
        self.declare_parameter('frcnn_min_score', 0.3)
        self.declare_parameter('frcnn_top_k', 2)
        self.bridge = CvBridge()
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        # Camera cache: subscriber callbacks only keep the latest data.
        self.latest_depth = None
        self.latest_depth_stamp = None
        self.camera_info = None
        self.latest_camera_info_stamp = None
        self.latest_stamp = None
        self.latest_rgb = None
        self.img_frame = "realsense_depth_frame"
        self.processing_frame = False
        self.sensor_callback_group = ReentrantCallbackGroup()
        self.timer_callback_group = MutuallyExclusiveCallbackGroup()
        self.tf_buffer = Buffer(cache_time=Duration(seconds=10))
        self.tf_listener = TransformListener(self.tf_buffer, self)


        # CLIP + Faster R-CNN
        self.clip_model, self.clip_preprocess = clip.load("ViT-B/32", device=self.device)
        self.frcnn_model = torchvision.models.detection.fasterrcnn_resnet50_fpn_v2(weights="DEFAULT").eval()

        # Labels and score tracking
        self.labels = ["Refrigerator","water dispenser", "sofa", "white toilet",  "office chair with wheels"]
        
        # Initialize score tracking - save directly to vlm_nav_pkg config
        # so semantic_nav reads the same file without manual copy
        try:
            nav_pkg_dir = get_package_share_directory('vlm_nav_pkg')
            self.score_records_file = os.path.join(
                nav_pkg_dir, 'config', 'example_object_detection_vlm.json'
            )
        except Exception:
            self.get_logger().warn(
                "vlm_nav_pkg not found, saving JSON to current directory"
            )
            self.score_records_file = "object_detection_vlm.json"
        reset_records = self.get_parameter(
            'reset_records_on_startup'
        ).get_parameter_value().bool_value
        if reset_records:
            self.score_records = {
                label: {"highest_score": 0.0, "last_detected": None, "position": None}
                for label in self.labels
            }
            self.save_score_records()
            self.get_logger().info(f"Score records reset: {self.score_records_file}")
        else:
            self.score_records = self.load_score_records()
            self.get_logger().info(f"Score records loaded: {self.score_records_file}")
        
        # Subscribers & Publishers
        self.pose_pub = self.create_publisher(PoseStamped, "/object_in_map", 10)

        self.rgb_sub = self.create_subscription(
            Image,
            "/intel_realsense_r200_rgb/image_raw",
            self.rgb_callback,
            10,
            callback_group=self.sensor_callback_group
        )
        self.depth_sub = self.create_subscription(
            Image,
            "/intel_realsense_r200_depth/depth/image_raw",
            self.depth_callback,
            10,
            callback_group=self.sensor_callback_group
        )
        self.camera_info_sub = self.create_subscription(
            CameraInfo, "/intel_realsense_r200_depth/camera_info",
            self.camera_info_callback, 10,
            callback_group=self.sensor_callback_group)
        self.annotated_pub = self.create_publisher(Image, "/percep/annotated_image", 10)
        interval_sec = self.get_parameter(
            'detection_interval_sec'
        ).get_parameter_value().double_value
        self.detection_timer = self.create_timer(
            interval_sec,
            self.process_frame,
            callback_group=self.timer_callback_group,
        )
        self.get_logger().info(
            f"Detection timer enabled: processing latest frame every {interval_sec:.2f}s"
        )


    def camera_info_callback(self, msg):
        self.camera_info = msg
        self.latest_camera_info_stamp = msg.header.stamp
        if msg.header.frame_id:
            self.img_frame = msg.header.frame_id

        

    def load_score_records(self):
        """Load existing score records or create new ones"""
        if os.path.exists(self.score_records_file):
            with open(self.score_records_file, 'r') as f:
                try:
                    return json.load(f)
                except json.JSONDecodeError:
                    self.get_logger().warn("Score file corrupted, creating new one")
        
        # Initialize with empty records
        return {label: {"highest_score": 0.0, "last_detected": None, "position":None} for label in self.labels}

    def save_score_records(self):
        """Save current score records to file"""
        with open(self.score_records_file, 'w') as f:
            json.dump(self.score_records, f, indent=4)

    def update_score_records(self, label, score, position=None):
        """Update records when confidence matches or exceeds the best score."""
        current_time = datetime.now().isoformat()
        
        if label not in self.score_records:
            self.score_records[label] = {
                "highest_score": score,
                "last_detected": current_time,
                "position": position
            }
            self.save_score_records()
            return True
        
        updated = False
        if score >= self.score_records[label]["highest_score"]:
            # Refresh the best score and associated position on ties or improvements.
            self.score_records[label]["highest_score"] = score
            if position is not None:
                self.score_records[label]["position"] = position
            updated = True
        
        # Always refresh timestamp (even without position/score change)
        self.score_records[label]["last_detected"] = current_time
        self.save_score_records()
        
        return updated


    def rgb_callback(self, msg):
        try:
            self.latest_rgb = self.bridge.imgmsg_to_cv2(msg, "bgr8")
            self.latest_stamp = msg.header.stamp
        except CvBridgeError as e:
            self.get_logger().error(f"CvBridge error: {e}")
            return


    def depth_callback(self, msg):
        try:
            self.latest_depth = self.bridge.imgmsg_to_cv2(msg, desired_encoding="32FC1")
            self.latest_depth_stamp = msg.header.stamp
        except CvBridgeError as e:
            self.get_logger().error(f"Depth CvBridge error: {e}")
            return

    def process_frame(self):
        if self.processing_frame:
            return

        if self.latest_rgb is None or self.latest_depth is None or self.camera_info is None:
            return

        if self.latest_stamp is None or self.latest_depth_stamp is None:
            return

        self.processing_frame = True
        start_time = time.perf_counter()
        try:
            cv_img = self.latest_rgb.copy()
            depth_img = self.latest_depth.copy()
            frame_stamp = self.latest_depth_stamp
            camera_info = self.camera_info
            image_tensor = transforms.ToTensor()(cv_img)
            frcnn_min_score = self.get_parameter(
                'frcnn_min_score'
            ).get_parameter_value().double_value
            frcnn_top_k = self.get_parameter(
                'frcnn_top_k'
            ).get_parameter_value().integer_value

            with torch.no_grad():
                frcnn_output = self.frcnn_model([image_tensor])[0]
            candidate_boxes = frcnn_output['boxes']
            candidate_scores = frcnn_output['scores']
            total_candidate_count = len(candidate_boxes)

            filtered_indices = [
                i for i, score in enumerate(candidate_scores)
                if float(score) >= frcnn_min_score
            ]
            filtered_indices.sort(key=lambda i: float(candidate_scores[i]), reverse=True)
            selected_indices = filtered_indices[:frcnn_top_k]
            candidate_boxes = [candidate_boxes[i] for i in selected_indices]
            self.get_logger().info(
                f"FRCNN candidates: total={total_candidate_count} "
                f"filtered={len(filtered_indices)} top_k={frcnn_top_k} "
                f"selected={len(candidate_boxes)}"
            )

            best_labels_per_box = {}  # key: box index, value: (label, score)

            for i, box in enumerate(candidate_boxes):
                max_score = -1
                best_label = None
                for label in self.labels:
                    _, score = self.match_label_box(label, image_tensor, [box])
                    if score > max_score:
                        max_score = score
                        best_label = label

                if max_score >= 0.25 and best_label is not None:
                    best_labels_per_box[i] = (best_label, max_score)
                    is_new_record = False
                    position = None

                    box_index = i
                    box = candidate_boxes[box_index]
                    x_min, y_min, x_max, y_max = map(int, box)
                    x_center = (x_min + x_max) // 2
                    y_center = (y_min + y_max) // 2
                    depth = float(depth_img[y_center, x_center])
                    if np.isnan(depth) or depth <= 0.0:
                        continue

                    fx = camera_info.k[0]
                    fy = camera_info.k[4]
                    cx = camera_info.k[2]
                    cy = camera_info.k[5]

                    X = (x_center - cx) * depth / fx
                    Y = (y_center - cy) * depth / fy
                    Z = depth

                    pose = PoseStamped()
                    pose.header.stamp = frame_stamp
                    pose.header.frame_id = self.img_frame
                    pose.pose.position.x = X
                    pose.pose.position.y = Y
                    pose.pose.position.z = Z
                    pose.pose.orientation.w = 1.0

                    # TF transform: camera frame -> map frame (semantic_nav 需要 map 坐标)
                    try:
                        ob2map = self.tf_buffer.transform(
                            pose, "map", timeout=Duration(seconds=1.0)
                        )

                        self.pose_pub.publish(ob2map)
                        self.get_logger().info(
                            f"object coordinates in map frame: x={ob2map.pose.position.x:.3f}, "
                            f"y={ob2map.pose.position.y:.3f}, z={ob2map.pose.position.z:.3f}"
                        )

                        position = {
                            "x": round(ob2map.pose.position.x, 4),
                            "y": round(ob2map.pose.position.y, 4),
                            "z": round(ob2map.pose.position.z, 4)
                        }
                        is_new_record = self.update_score_records(best_label, max_score, position=position)
                    except Exception as e:
                        self.get_logger().error(f"TF transform to map failed: {e}")
                        is_new_record = False

                    log_msg = f"{best_label} detected with confidence score {max_score:.2f}"
                    if is_new_record:
                        log_msg += " (NEW RECORD!)"
                    self.get_logger().info(log_msg)

            matched_boxes = [candidate_boxes[i] for i in best_labels_per_box.keys()]
            matched_labels = [f"{lbl} ({score:.2f})" for lbl, score in best_labels_per_box.values()]

            annotated_img = self.draw_boxes(cv_img, matched_boxes, matched_labels)
            try:
                self.annotated_pub.publish(self.bridge.cv2_to_imgmsg(annotated_img, encoding="bgr8"))
            except CvBridgeError as e:
                self.get_logger().error(f"Publish error: {e}")
        finally:
            elapsed_ms = (time.perf_counter() - start_time) * 1000.0
            self.get_logger().info(f"Detection cycle took {elapsed_ms:.1f} ms")
            self.processing_frame = False
        
    def match_label_box(self, label, image_tensor, boxes):
        max_sim, best_box = -1, None
        tokenized_text = clip.tokenize([label]).to(self.device)
        text_feat = F.normalize(self.clip_model.encode_text(tokenized_text), dim=-1)

        for box in boxes:
            x_min, y_min, x_max, y_max = map(int, box)
            crop = transforms.functional.to_pil_image(image_tensor[:, y_min:y_max, x_min:x_max])
            img_feat = F.normalize(self.clip_model.encode_image(self.clip_preprocess(crop).unsqueeze(0).to(self.device)), dim=-1)
            sim = F.cosine_similarity(img_feat, text_feat).item()
            if sim > max_sim:
                max_sim, best_box = sim, box
        return best_box, max_sim

    def draw_boxes(self, img, boxes, labels):
        for i, box in enumerate(boxes):
            x_min, y_min, x_max, y_max = map(int, box)
            cv2.rectangle(img, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
            if labels:
                cv2.putText(img, labels[i], (x_min, y_min-10), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 0, 0), 2)
        return img

def main(args=None):
    rclpy.init(args=args)
    node = ClipDetectionNode()
    executor = MultiThreadedExecutor(num_threads=2)
    executor.add_node(node)
    try:
        executor.spin()
    except KeyboardInterrupt:
        node.get_logger().info("KeyboardInterrupt received, shutting down...")
    finally:
        node.save_score_records()
        executor.shutdown()
        node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()

if __name__ == "__main__":
    main()
