#!/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
from math import isfinite
import numpy as np
import yaml
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"

        # Behavior controls for semantic memory persistence and position quality.
        self.declare_parameter('reset_records_on_startup', False)
        self.declare_parameter('position_abs_limit_xy', 50.0)
        self.declare_parameter('position_z_min', -2.0)
        self.declare_parameter('position_z_max', 5.0)
        self.declare_parameter('camera_frame_fallback', 'realsense_depth_frame')
        self.declare_parameter(
            'map_yaml_path',
            os.environ.get('NAV2_MAP_PATH', '')
        )
        self.declare_parameter('best_position_confirm_frames', 3)
        self.declare_parameter('best_position_confirm_radius', 0.8)
        self.reset_records_on_startup = bool(
            self.get_parameter('reset_records_on_startup').value
        )
        self.position_abs_limit_xy = float(
            self.get_parameter('position_abs_limit_xy').value
        )
        self.position_z_min = float(self.get_parameter('position_z_min').value)
        self.position_z_max = float(self.get_parameter('position_z_max').value)
        self.camera_frame_fallback = str(
            self.get_parameter('camera_frame_fallback').value
        )
        self.map_yaml_path = str(self.get_parameter('map_yaml_path').value).strip()
        self.best_position_confirm_frames = int(
            self.get_parameter('best_position_confirm_frames').value
        )
        self.best_position_confirm_radius = float(
            self.get_parameter('best_position_confirm_radius').value
        )
        if not self.camera_frame_fallback:
            self.camera_frame_fallback = 'realsense_depth_frame'
        if self.best_position_confirm_frames < 1:
            self.best_position_confirm_frames = 1
        if self.best_position_confirm_radius <= 0.0:
            self.best_position_confirm_radius = 0.8

        #camera init
        self.latest_depth = None
        self.camera_info = None
        self.latest_stamp = None
        self.latest_rgb = None
        self.latest_depth_stamp = None
        self.img_frame = self.camera_frame_fallback
        self.last_tf_warn = ""
        self.last_frame_warn = ""
        # Keep a longer TF history to tolerate localization TF gaps.
        self.tf_buffer = Buffer(cache_time=Duration(seconds=120))
        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"]
        self.best_position_candidates = {label: [] for label in self.labels}
        self.map_bounds = self._load_map_bounds()
        
        # 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 = os.path.join(
                nav_pkg_dir, 'config', 'object_detection_vlm.json'
            )
        if self.reset_records_on_startup:
            self.score_records = {
                label: self._empty_record()
                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()
            # Persist normalized schema for legacy files.
            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
        incoming_frame = (msg.header.frame_id or "").strip()
        if not incoming_frame:
            return
        if incoming_frame == self.img_frame:
            return

        # Gz bridge may publish scoped names like
        # burger/realsense_link/intel_realsense_r200_depth which are not in TF.
        # Keep the ROS TF frame fallback in that case.
        if "/" in incoming_frame:
            warn_msg = (
                f"Ignoring scoped camera frame '{incoming_frame}', "
                f"using fallback '{self.camera_frame_fallback}'"
            )
            if warn_msg != self.last_frame_warn:
                self.get_logger().warn(warn_msg)
                self.last_frame_warn = warn_msg
            self.img_frame = self.camera_frame_fallback
            return

        self.img_frame = incoming_frame
        self.get_logger().info(f"Camera info frame updated: {self.img_frame}")

        

    def _empty_record(self):
        return {
            "highest_score": 0.0,
            "last_detected": None,
            # Backward-compatible field consumed by semantic_nav.
            "position": None,
            # New fields:
            "best_position": None,
            "last_position": None,
        }

    def _resolve_map_yaml_path(self):
        path = os.path.expanduser(self.map_yaml_path) if self.map_yaml_path else ''
        if path and os.path.isfile(path):
            return path
        try:
            tb3_dir = get_package_share_directory('turtlebot3_gazebo')
            default_path = os.path.join(tb3_dir, 'map', 'office_map.yaml')
            if os.path.isfile(default_path):
                return default_path
        except Exception:
            pass
        return ''

    def _load_map_bounds(self):
        map_yaml = self._resolve_map_yaml_path()
        if not map_yaml:
            self.get_logger().warn(
                "Map bounds disabled: map_yaml_path not found. "
                "Set NAV2_MAP_PATH or parameter map_yaml_path."
            )
            return None
        try:
            with open(map_yaml, 'r') as f:
                map_cfg = yaml.safe_load(f)
            resolution = float(map_cfg['resolution'])
            origin = map_cfg['origin']
            origin_x = float(origin[0])
            origin_y = float(origin[1])
            image_rel = str(map_cfg['image'])
            image_path = image_rel
            if not os.path.isabs(image_path):
                image_path = os.path.join(os.path.dirname(map_yaml), image_rel)
            map_img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
            if map_img is None:
                raise RuntimeError(f"cannot open map image: {image_path}")
            height, width = map_img.shape[:2]
            bounds = {
                'min_x': origin_x,
                'max_x': origin_x + width * resolution,
                'min_y': origin_y,
                'max_y': origin_y + height * resolution,
                'map_yaml': map_yaml,
            }
            self.get_logger().info(
                "Map bounds loaded: "
                f"x=[{bounds['min_x']:.3f}, {bounds['max_x']:.3f}], "
                f"y=[{bounds['min_y']:.3f}, {bounds['max_y']:.3f}]"
            )
            return bounds
        except Exception as e:
            self.get_logger().warn(
                f"Map bounds disabled due to load error: {e}"
            )
            return None

    def _is_in_map_bounds(self, x, y):
        if self.map_bounds is None:
            return True
        return (
            self.map_bounds['min_x'] <= x <= self.map_bounds['max_x']
            and self.map_bounds['min_y'] <= y <= self.map_bounds['max_y']
        )

    def _positions_consistent(self, positions):
        if len(positions) < self.best_position_confirm_frames:
            return False
        center_x = float(np.median([p['x'] for p in positions]))
        center_y = float(np.median([p['y'] for p in positions]))
        max_dist = 0.0
        for p in positions:
            dx = p['x'] - center_x
            dy = p['y'] - center_y
            dist = float(np.hypot(dx, dy))
            if dist > max_dist:
                max_dist = dist
        return max_dist <= self.best_position_confirm_radius

    def _median_position(self, positions):
        if not positions:
            return None
        return {
            "x": round(float(np.median([p['x'] for p in positions])), 4),
            "y": round(float(np.median([p['y'] for p in positions])), 4),
            "z": round(float(np.median([p['z'] for p in positions])), 4),
        }

    def _sanitize_position(self, position):
        """Return normalized position dict or None if invalid."""
        if position is None:
            return None
        try:
            x = float(position["x"])
            y = float(position["y"])
            z = float(position.get("z", 0.0))
        except Exception:
            return None

        if not (isfinite(x) and isfinite(y) and isfinite(z)):
            return None
        if abs(x) > self.position_abs_limit_xy or abs(y) > self.position_abs_limit_xy:
            return None
        if z < self.position_z_min or z > self.position_z_max:
            return None
        if not self._is_in_map_bounds(x, y):
            return None

        return {
            "x": round(x, 4),
            "y": round(y, 4),
            "z": round(z, 4),
        }

    def load_score_records(self):
        """Load existing score records and normalize to current schema."""
        raw_records = {}
        if os.path.exists(self.score_records_file):
            with open(self.score_records_file, 'r') as f:
                try:
                    raw_records = json.load(f)
                except json.JSONDecodeError:
                    self.get_logger().warn("Score file corrupted, creating new one")
                    raw_records = {}
        if not isinstance(raw_records, dict):
            self.get_logger().warn("Score file format invalid, recreating records")
            raw_records = {}

        normalized = {}
        for label in self.labels:
            src = raw_records.get(label, {})
            dst = self._empty_record()

            highest_score = src.get("highest_score", 0.0)
            try:
                highest_score = float(highest_score)
            except Exception:
                highest_score = 0.0
            if not isfinite(highest_score):
                highest_score = 0.0
            dst["highest_score"] = max(0.0, highest_score)
            dst["last_detected"] = src.get("last_detected")

            legacy_pos = self._sanitize_position(src.get("position"))
            best_pos = self._sanitize_position(src.get("best_position"))
            last_pos = self._sanitize_position(src.get("last_position"))

            if best_pos is None:
                best_pos = legacy_pos
            if last_pos is None:
                last_pos = legacy_pos

            dst["best_position"] = best_pos
            dst["last_position"] = last_pos
            dst["position"] = best_pos
            normalized[label] = dst

        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):
        """Update semantic memory with dual-track positions.

        - best_position: saved when confidence reaches a new high score.
        - last_position: latest valid position, regardless of score.
        - position: backward-compatible mirror of best_position.
        """
        current_time = datetime.now().isoformat()

        if label not in self.score_records:
            self.score_records[label] = self._empty_record()

        record = self.score_records[label]
        changed = False
        new_best = False

        # Always refresh timestamp when object is observed.
        record["last_detected"] = current_time

        try:
            score = float(score)
        except Exception:
            score = 0.0
        if not isfinite(score):
            score = 0.0

        normalized_position = self._sanitize_position(position)
        if normalized_position is not None:
            if record.get("last_position") != normalized_position:
                record["last_position"] = normalized_position
                changed = True

            candidates = self.best_position_candidates.setdefault(label, [])
            candidates.append({
                "position": normalized_position,
                "score": score,
            })
            if len(candidates) > self.best_position_confirm_frames:
                candidates.pop(0)

            highest_score = float(record.get("highest_score", 0.0))
            recent_positions = [c["position"] for c in candidates]
            if (
                len(candidates) >= self.best_position_confirm_frames
                and self._positions_consistent(recent_positions)
            ):
                candidate_score = max(float(c["score"]) for c in candidates)
                if candidate_score > highest_score:
                    median_pos = self._median_position(recent_positions)
                    record["highest_score"] = candidate_score
                    record["best_position"] = median_pos
                    record["position"] = median_pos
                    changed = True
                    new_best = True
            elif record.get("position") != record.get("best_position"):
                record["position"] = record.get("best_position")
                changed = True
        else:
            if position is not None:
                self.get_logger().warn(
                    f"Invalid map position dropped for '{label}': {position}"
                )
            if record.get("position") != record.get("best_position"):
                record["position"] = record.get("best_position")
                changed = True

        if changed:
            self.save_score_records()

        return new_best


    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 is_new_record as False first
                is_new_record = 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

                    pose = PoseStamped()
                    pose.header.stamp = Time().to_msg()  # 零时间=使用最新可用TF
                    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)
                        )
                    except Exception as e:
                        err_msg = str(e)
                        # Fallback path:
                        # evaluate source->odom and map->odom independently at latest time.
                        # This avoids "latest common time" failures when map->odom is sparse.
                        if "source_frame does not exist" in err_msg:
                            try:
                                pose.header.frame_id = self.camera_frame_fallback
                                ob2map = self.tf_buffer.transform(
                                    pose, "map", timeout=Duration(seconds=0.2)
                                )
                                warn_msg = (
                                    "Recovered TF by forcing fallback frame: "
                                    f"{self.camera_frame_fallback}"
                                )
                                if warn_msg != self.last_tf_warn:
                                    self.get_logger().warn(warn_msg)
                                    self.last_tf_warn = warn_msg
                            except Exception as fallback_err:
                                self.get_logger().error(
                                    f"TF transform to map failed: {fallback_err}"
                                )
                                is_new_record = False
                                continue
                        elif "extrapolation into the past" in err_msg:
                            try:
                                ob2map = self.tf_buffer.transform_full(
                                    pose,
                                    "map",
                                    Time(),
                                    "odom",
                                    timeout=Duration(seconds=0.2)
                                )
                                warn_msg = (
                                    "Recovered TF with transform_full fallback via fixed frame 'odom'"
                                )
                                if warn_msg != self.last_tf_warn:
                                    self.get_logger().warn(warn_msg)
                                    self.last_tf_warn = warn_msg
                            except Exception as fallback_err:
                                self.get_logger().error(
                                    f"TF transform to map failed: {fallback_err}"
                                )
                                is_new_record = False
                                continue
                        else:
                            self.get_logger().error(f"TF transform to map failed: {e}")
                            is_new_record = False
                            continue

                    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": ob2map.pose.position.x,
                        "y": ob2map.pose.position.y,
                        "z": ob2map.pose.position.z
                    }
                    is_new_record = self.update_score_records(best_label, max_score, position=position)


                # Logging
                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}")

        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()
