#!/usr/bin/env python3
import json
import clip
import torch
import os
from math import isfinite
import yaml
import cv2
import rclpy
from rclpy.node import Node
from geometry_msgs.msg import PoseStamped
from nav2_simple_commander.robot_navigator import BasicNavigator, TaskResult
from ament_index_python.packages import get_package_share_directory
class SemanticNavNode(Node):
    def __init__(self, json_path: str):
        super().__init__('semantic_nav_node')

        # Load object library
        with open(json_path) as f:
            self.object_lib = json.load(f)
        self.map_bounds = self._load_map_bounds()

        # Load CLIP
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model, _ = clip.load("ViT-B/32", device=self.device)

        # Build text embeddings for all object names
        self.object_names = list(self.object_lib.keys())
        text_tokens = clip.tokenize(self.object_names).to(self.device)
        with torch.no_grad():
            text_features = self.model.encode_text(text_tokens)
            self.text_features = text_features / text_features.norm(dim=-1, keepdim=True)

        # Initialize navigator
        self.navigator = BasicNavigator()
        self.navigator.waitUntilNav2Active()
        self.get_logger().info("Navigator ready")

    def _load_map_bounds(self):
        map_yaml = os.environ.get('NAV2_MAP_PATH', '').strip()
        map_yaml = os.path.expanduser(map_yaml) if map_yaml else ''
        if not map_yaml or not os.path.isfile(map_yaml):
            try:
                tb3_dir = get_package_share_directory('turtlebot3_gazebo')
                map_yaml = os.path.join(tb3_dir, 'map', 'office_map.yaml')
            except Exception:
                map_yaml = ''
        if not map_yaml or not os.path.isfile(map_yaml):
            self.get_logger().warn(
                "Map bounds unavailable in semantic_nav; "
                "fallback selection will skip bounds checks."
            )
            return None
        try:
            with open(map_yaml, 'r') as f:
                cfg = yaml.safe_load(f)
            resolution = float(cfg['resolution'])
            ox, oy = float(cfg['origin'][0]), float(cfg['origin'][1])
            image_path = str(cfg['image'])
            if not os.path.isabs(image_path):
                image_path = os.path.join(os.path.dirname(map_yaml), image_path)
            img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
            if img is None:
                raise RuntimeError(f"cannot open map image: {image_path}")
            h, w = img.shape[:2]
            bounds = {
                "min_x": ox,
                "max_x": ox + w * resolution,
                "min_y": oy,
                "max_y": oy + h * resolution,
            }
            self.get_logger().info(
                "Semantic map bounds: "
                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"Failed to load map bounds: {e}")
            return None

    def _sanitize_position(self, pos):
        if not isinstance(pos, dict):
            return None
        try:
            x = float(pos["x"])
            y = float(pos["y"])
            z = float(pos.get("z", 0.0))
        except Exception:
            return None
        if not (isfinite(x) and isfinite(y) and isfinite(z)):
            return None
        return {"x": x, "y": y, "z": z}

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

    def query_object(self, prompt: str):
        tokens = clip.tokenize([prompt]).to(self.device)
        with torch.no_grad():
            query_feat = self.model.encode_text(tokens)
            query_feat /= query_feat.norm(dim=-1, keepdim=True)

        sims = (query_feat @ self.text_features.T).squeeze(0)
        best_idx = sims.argmax().item()
        best_name = self.object_names[best_idx]

        # Align with vlm_detection:
        # prefer best_position/position, fallback to last_position when best is out-of-bounds.
        obj_info = self.object_lib.get(best_name, {})
        best_pos_raw = obj_info.get("best_position")
        if best_pos_raw is None:
            best_pos_raw = obj_info.get("position")
        last_pos_raw = obj_info.get("last_position")
        best_pos = self._sanitize_position(best_pos_raw)
        last_pos = self._sanitize_position(last_pos_raw)

        if best_pos is not None and self._in_map_bounds(best_pos):
            return best_name, best_pos

        if best_pos is not None and not self._in_map_bounds(best_pos):
            self.get_logger().warn(
                f"'{best_name}' best_position out of map bounds: {best_pos}"
            )
            if last_pos is not None and self._in_map_bounds(last_pos):
                self.get_logger().warn(
                    f"Falling back to last_position for '{best_name}': {last_pos}"
                )
                return best_name, last_pos

        if last_pos is not None and self._in_map_bounds(last_pos):
            self.get_logger().warn(
                f"Using last_position for '{best_name}': {last_pos}"
            )
            return best_name, last_pos

        if best_pos is None and last_pos is None:
            self.get_logger().warn(
                f"'{best_name}' matched but has no usable position in JSON. "
                f"Run vlm_detection in AMCL mode first!"
            )
            return best_name, None

        self.get_logger().warn(
            f"'{best_name}' has only out-of-bounds position(s), cannot navigate."
        )
        return best_name, None

    def navigate_to_object(self, prompt: str):
        name, pos = self.query_object(prompt)
        if pos is None:
            self.get_logger().warn(f"Cannot navigate: no valid position for '{name}'. Please run vlm_detection first!")
            return

        goal_pose = PoseStamped()
        goal_pose.header.frame_id = "map"
        goal_pose.header.stamp = self.navigator.get_clock().now().to_msg()
        goal_pose.pose.position.x = pos["x"]
        goal_pose.pose.position.y = pos["y"]
        goal_pose.pose.position.z = pos.get("z", 0.0)
        goal_pose.pose.orientation.z = 0.0
        goal_pose.pose.orientation.w = 1.0

        self.navigator.followWaypoints([goal_pose])
        self.get_logger().info(f"Navigating to {name} at {pos}")

        # Wait until navigation completes
        while not self.navigator.isTaskComplete():
            feedback = self.navigator.getFeedback()
            if feedback:
                self.get_logger().info(
                    f"Executing waypoint {feedback.current_waypoint + 1}/1"
                )

        result = self.navigator.getResult()
        if result == TaskResult.SUCCEEDED:
            self.get_logger().info("Navigation succeeded")
        elif result == TaskResult.CANCELED:
            self.get_logger().info("Navigation canceled")
        elif result == TaskResult.FAILED:
            self.get_logger().info("Navigation failed")
            
            



def main():
    rclpy.init()
    package_share_dir = get_package_share_directory('vlm_nav_pkg')
    json_file_path = os.path.join(package_share_dir, 'config', 'example_object_detection_vlm.json')
    node = SemanticNavNode(json_file_path)

    user_input = input("Where do you want to go: ")
    node.navigate_to_object(user_input)

    rclpy.shutdown()


if __name__ == "__main__":
    main()
