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




    

class ClipDetectionNode(Node):
    def __init__(self):
        super().__init__('clip_detector')
        self.bridge = CvBridge()
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        #camera init
        self.latest_depth = None
        self.camera_info = None
        self.latest_stamp = None
        self.latest_depth_stamp = None
        self.last_tf_warn_time = None
        self.latest_rgb = None
        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"]

        # Save directly to vlm_nav_pkg config so semantic_nav reads the same file.
        try:
            nav_pkg_dir = get_package_share_directory('vlm_nav_pkg')
            default_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"
            )
            default_records_file = "object_detection_vlm.json"

        self.declare_parameter("score_records_file", default_records_file)
        self.declare_parameter("reset_records_on_startup", False)
        configured_file = self.get_parameter("score_records_file").value
        self.score_records_file = configured_file if configured_file else default_records_file
        self.reset_records_on_startup = self.get_parameter("reset_records_on_startup").value

        if self.reset_records_on_startup:
            self.score_records = self.get_empty_score_records()
            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.save_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
        )
        self.depth_sub = self.create_subscription(
            Image,
            "/intel_realsense_r200_depth/depth/image_raw",
            self.depth_callback,
            10
        )
        self.camera_info_sub = self.create_subscription(
            CameraInfo, "/intel_realsense_r200_depth/camera_info",
            self.camera_info_callback, 10)
        self.annotated_pub = self.create_publisher(Image, "/percep/annotated_image", 10)


    def camera_info_callback(self, msg):
        self.camera_info = msg
        self.img_frame = "realsense_depth_frame"
        self.get_logger().info(f"Camera info received, using frame: {self.img_frame}")

        

    def get_empty_score_records(self):
        return {
            label: {"highest_score": 0.0, "last_detected": None, "position": None}
            for label in self.labels
        }

    def load_score_records(self):
        """Load existing score records and normalize missing labels/fields."""
        empty_records = self.get_empty_score_records()
        if not os.path.exists(self.score_records_file):
            self.get_logger().warn(
                f"Score records file not found, creating new one: {self.score_records_file}"
            )
            return empty_records

        with open(self.score_records_file, 'r') as f:
            try:
                loaded = json.load(f)
            except json.JSONDecodeError:
                self.get_logger().warn("Score file corrupted, creating new one")
                return empty_records

        if not isinstance(loaded, dict):
            self.get_logger().warn("Score file format is invalid, creating new one")
            return empty_records

        def normalize_record(record):
            if not isinstance(record, dict):
                record = {}

            highest_score = record.get("highest_score", 0.0)
            try:
                highest_score = float(highest_score)
            except (TypeError, ValueError):
                highest_score = 0.0
            highest_score = max(0.0, highest_score)

            last_detected = record.get("last_detected")
            position = record.get("position")
            if position is not None and not isinstance(position, dict):
                position = None

            return {
                "highest_score": highest_score,
                "last_detected": last_detected,
                "position": position
            }

        normalized = {}
        for label, record in loaded.items():
            normalized[label] = normalize_record(record)

        for label in self.labels:
            if label not in normalized:
                normalized[label] = empty_records[label]

        return normalized

    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):
        """Keep highest score per label; equal score is treated as an update."""
        current_time = datetime.now().isoformat()

        if label not in self.score_records:
            self.score_records[label] = {
                "highest_score": 0.0,
                "last_detected": None,
                "position": None
            }

        record = self.score_records[label]
        updated = False
        if score >= record["highest_score"]:
            # Equal score is also accepted so records can refresh over time.
            record["highest_score"] = score
            if position is not None:
                record["position"] = position
            updated = True

        record["last_detected"] = current_time
        if updated:
            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  # <-- Correct assignment
        except CvBridgeError as e:
            self.get_logger().error(f"CvBridge error: {e}")
            return
        self.process_frame()


    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
        self.process_frame()

    def process_frame(self):
        if self.latest_rgb is None or self.latest_depth is None:
            return

        cv_img = self.latest_rgb.copy()
        image_tensor = transforms.ToTensor()(cv_img)

        with torch.no_grad():
            frcnn_output = self.frcnn_model([image_tensor])[0]
        candidate_boxes = frcnn_output['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)
                # Initialize record_updated as False first
                record_updated = False
                position = None

                if self.latest_depth is not None and self.camera_info is not 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(self.latest_depth[y_center, x_center])
                    if np.isnan(depth) or depth <= 0.0:
                        continue

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

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

                    stamp_msg = self.latest_depth_stamp if self.latest_depth_stamp is not None else self.latest_stamp
                    if stamp_msg is None:
                        continue
                    query_time = Time.from_msg(
                        stamp_msg, clock_type=self.get_clock().clock_type
                    )
                    tf_stamp_msg = stamp_msg
                    has_tf_at_stamp = self.tf_buffer.can_transform(
                        "map", self.img_frame, query_time, timeout=Duration(seconds=0.05)
                    )
                    if not has_tf_at_stamp:
                        latest_query_time = Time(clock_type=self.get_clock().clock_type)
                        has_tf_latest = self.tf_buffer.can_transform(
                            "map", self.img_frame, latest_query_time, timeout=Duration(seconds=0.05)
                        )
                        if not has_tf_latest:
                            now_time = self.get_clock().now()
                            if (
                                self.last_tf_warn_time is None
                                or (now_time - self.last_tf_warn_time).nanoseconds > 2_000_000_000
                            ):
                                self.get_logger().warn(
                                    "TF unavailable at sensor stamp and latest time; skip this detection frame."
                                )
                                self.last_tf_warn_time = now_time
                            continue
                        tf_stamp_msg = latest_query_time.to_msg()

                    pose = PoseStamped()
                    pose.header.stamp = tf_stamp_msg
                    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=0.05)
                        )

                        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)
                        }
                        record_updated = 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}")
                        record_updated = False


                # Logging
                log_msg = f"{best_label} detected with confidence score {max_score:.2f}"
                if record_updated:
                    log_msg += " (BEST SCORE RECORD UPDATED)"
                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}")

        self.latest_rgb, self.latest_depth = None, None
        self.latest_stamp = None  # 同步清空时间戳，避免TF查询残留旧stamp
        self.latest_depth_stamp = None
        
    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()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        node.get_logger().info("KeyboardInterrupt received, shutting down...")
    finally:
        node.save_score_records()
        node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()

if __name__ == "__main__":
    main()
