#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
import torch
import numpy as np
import time
import math
import sys
import os

# 路径修复
current_dir = os.path.dirname(os.path.abspath(__file__))
package_root = os.path.abspath(os.path.join(current_dir, "../../../"))
if package_root not in sys.path:
    sys.path.append(package_root)

from soft_arm_sim.deeplearning.models.mlp_network import SimpleIKNet
from soft_arm_sim.deeplearning.layers.differentiable_pcc import DifferentiablePCC

# 引入 MarkerArray
from visualization_msgs.msg import Marker, MarkerArray
from geometry_msgs.msg import Point

class InferenceNode(Node):
    def __init__(self):
        super().__init__('inference_node')
        self.get_logger().info("Initializing PINN Inference Node (MarkerArray Mode)...")
        
        # 1. 自动寻找模型路径
        node_dir = os.path.dirname(os.path.abspath(__file__))
        self.model_path = os.path.join(node_dir, "../checkpoints/pinn_basic.pth")
        
        self.device = torch.device("cpu")
        
        # 2. 初始化网络和FK层
        self.model = SimpleIKNet(input_dim=3, output_dim=6).to(self.device)
        self.fk_layer = DifferentiablePCC().to(self.device) # 用于计算显示的坐标
        
        self.load_model()
        
        # 3. 发布者
        # 改用 MarkerArray 发布整个机械臂形态
        self.viz_pub = self.create_publisher(MarkerArray, 'soft_arm_viz', 10)
        self.marker_pub = self.create_publisher(Marker, 'target_marker', 10)
        
        self.timer = self.create_timer(0.02, self.control_loop)
        self.start_time = time.time()

    def load_model(self):
        try:
            self.model.load_state_dict(torch.load(self.model_path, map_location=self.device))
            self.model.eval()
            self.get_logger().info(f"Loaded model from: {self.model_path}")
        except Exception as e:
            self.get_logger().error(f"Error loading model: {e}")

    def control_loop(self):
        # 1. 生成画圆轨迹
        t = time.time() - self.start_time
        radius = 0.15
        center_z = 0.50
        freq = 1.0
        
        target_pos = [
            radius * math.cos(freq * t),
            radius * math.sin(freq * t),
            center_z
        ]

        # 2. 推理
        input_tensor = torch.tensor([target_pos], dtype=torch.float32).to(self.device)
        with torch.no_grad():
            pred_joints = self.model(input_tensor) # 得到角度
            
            # 3. 通过 FK 层计算骨架坐标 (用于可视化)
            # pred_joints: (1, 6)
            # points shape: (1, N_Points, 3)
            _, points_tensor = self.fk_layer(pred_joints)
            
        # 转为 numpy 用于 ROS 消息
        # points_np shape: (N_Points, 3)
        points_np = points_tensor.cpu().numpy()[0]
        
        # 4. 发布可视化
        self.publish_arm_viz(points_np)
        self.publish_target(target_pos)

    def publish_arm_viz(self, points):
        """
        根据骨架点生成 MarkerArray
        points: List of [x, y, z]
        """
        array_msg = MarkerArray()
        timestamp = self.get_clock().now().to_msg()
        
        # --- A. 骨架连线 (LINE_STRIP) ---
        line_marker = Marker()
        line_marker.header.frame_id = "base_link" # 确保 Rviz 里 Fixed Frame 也是这个
        line_marker.header.stamp = timestamp
        line_marker.ns = "backbone"
        line_marker.id = 0
        line_marker.type = Marker.LINE_STRIP
        line_marker.action = Marker.ADD
        line_marker.scale.x = 0.02 # 线条粗细
        line_marker.color.a = 1.0
        line_marker.color.r = 0.0
        line_marker.color.g = 0.8
        line_marker.color.b = 1.0 # 青色
        
        for p in points:
            pt = Point()
            pt.x, pt.y, pt.z = float(p[0]), float(p[1]), float(p[2])
            line_marker.points.append(pt)
            
        array_msg.markers.append(line_marker)
        
        # --- B. 关节点/圆盘 (SPHERE) ---
        for i, p in enumerate(points):
            disk_marker = Marker()
            disk_marker.header.frame_id = "base_link"
            disk_marker.header.stamp = timestamp
            disk_marker.ns = "disks"
            disk_marker.id = i + 1 # ID 不能重复
            disk_marker.type = Marker.SPHERE
            disk_marker.action = Marker.ADD
            
            disk_marker.pose.position.x = float(p[0])
            disk_marker.pose.position.y = float(p[1])
            disk_marker.pose.position.z = float(p[2])
            
            # 只有中间的点代表磁盘，原点和末端可以小一点
            disk_marker.scale.x = 0.08 
            disk_marker.scale.y = 0.08 
            disk_marker.scale.z = 0.02 # 扁的圆盘
            
            disk_marker.color.a = 1.0
            disk_marker.color.r = 1.0
            disk_marker.color.g = 0.6
            disk_marker.color.b = 0.0 # 橙色
            
            array_msg.markers.append(disk_marker)
            
        self.viz_pub.publish(array_msg)

    def publish_target(self, pos):
        marker = Marker()
        marker.header.frame_id = "base_link"
        marker.header.stamp = self.get_clock().now().to_msg()
        marker.ns = "target"
        marker.id = 999
        marker.type = Marker.SPHERE
        marker.action = Marker.ADD
        marker.pose.position.x, marker.pose.position.y, marker.pose.position.z = pos
        marker.scale.x = marker.scale.y = marker.scale.z = 0.05
        marker.color.a = 1.0; marker.color.r = 1.0 # 红球
        self.marker_pub.publish(marker)

def main(args=None):
    rclpy.init(args=args)
    node = InferenceNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()