import rclpy
from rclpy.node import Node
from geometry_msgs.msg import PoseStamped, Point
from std_msgs.msg import Float64MultiArray
from visualization_msgs.msg import Marker
import torch
import numpy as np
import time
import threading
import os
import sys

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

from pinn_model import PINN_IK
from soft_arm_sim.model.pcc_kinematics import SoftArmKinematics

class PinnController(Node):
    def __init__(self):
        super().__init__('pinn_controller')
        
        # 1. 加载模型和统计数据
        self.device = torch.device("cpu")
        self.model = PINN_IK().to(self.device)
        self.stats = None # 用于存储归一化参数
        
        model_path = os.path.join(current_dir, "best_model_v2.pth")
        
        if os.path.exists(model_path):
            checkpoint = torch.load(model_path, map_location=self.device)
            
            # 加载权重
            self.model.load_state_dict(checkpoint['model_state'])
            # 加载归一化参数
            self.stats = checkpoint['stats']
            
            self.model.eval()
            self.get_logger().info(f"✅ 模型 V2 (带归一化) 加载成功")
        else:
            self.get_logger().error(f"❌ 找不到模型: {model_path}，请运行新的 train.py")

        self.fk_solver = SoftArmKinematics(
            num_sections=3, section_length=0.24, disks_per_section=3, disk_radius=0.033
        )

        self.cmd_pub = self.create_publisher(Float64MultiArray, 'soft_arm/command', 10)
        self.traj_pub = self.create_publisher(Marker, 'planned_trajectory', 10)
        self.target_marker_pub = self.create_publisher(Marker, 'target_marker', 10)
        self.create_subscription(PoseStamped, '/goal_pose', self.goal_callback, 10)
        self.current_pos = np.array([0.0, 0.0, 0.72]) 
        self.get_logger().info("等待目标指令...")

    def goal_callback(self, msg):
        target_pos = np.array([msg.pose.position.x, msg.pose.position.y, 0.5]) # 强制Z=0.5测试
        
        dist = np.linalg.norm(target_pos)
        if dist > 0.72:
            self.get_logger().warn(f"⚠️ 目标点太远 ({dist:.2f}m)，可能无法到达")
        
        self.get_logger().info(f"收到目标: {target_pos}，开始规划...")
        self.visualize_target(target_pos)
        threading.Thread(target=self.execute_move, args=(target_pos,)).start()

    def preprocess_input(self, pos):
        """ 将真实位置 (x,y,z) 归一化 """
        if self.stats is None: return pos
        X_mean = self.stats['X_mean']
        X_std = self.stats['X_std']
        return (pos - X_mean) / X_std

    def postprocess_output(self, norm_angles):
        """ 将归一化角度还原为真实角度 """
        if self.stats is None: return norm_angles
        y_mean = self.stats['y_mean']
        y_std = self.stats['y_std']
        return norm_angles * y_std + y_mean

    def execute_move(self, target_pos):
        steps = 50
        trajectory = []
        for t in np.linspace(0, 1, steps):
            pt = self.current_pos + t * (target_pos - self.current_pos)
            trajectory.append(pt)
        self.visualize_trajectory(trajectory)

        final_angles = None
        for pt in trajectory:
            with torch.no_grad():
                # 1. 预处理 (归一化)
                norm_input = self.preprocess_input(pt)
                tensor_input = torch.FloatTensor(norm_input).to(self.device)
                
                # 2. 推理
                norm_output = self.model(tensor_input).numpy()
                
                # 3. 后处理 (反归一化)
                real_angles = self.postprocess_output(norm_output)
                final_angles = real_angles
            
            msg = Float64MultiArray()
            msg.data = real_angles.tolist()
            self.cmd_pub.publish(msg)
            time.sleep(0.04)
            
        self.current_pos = target_pos
        self.get_logger().info("运动完成")
        if final_angles is not None:
            self.check_accuracy(target_pos, final_angles)

    def check_accuracy(self, target_pos, angles):
        q_input = []
        for i in range(3):
            q_input.append((angles[i*2], angles[i*2+1], 0.24))
        transforms, _ = self.fk_solver.forward(q_input)
        actual_pos = transforms[-1][:3, 3]
        error = np.linalg.norm(target_pos - actual_pos)
        
        print("\n" + "="*30)
        print(f"🎯 目标位置: {target_pos}")
        print(f"🤖 实际位置: {actual_pos}")
        print(f"❌ 误差: {error:.4f} m")
        print("="*30 + "\n")
    
    # ... visualize 函数保持不变 ...
    def visualize_target(self, pos):
        marker = Marker()
        marker.header.frame_id = "base_link"
        marker.type = Marker.SPHERE; marker.action = Marker.ADD
        marker.scale.x = 0.05; marker.scale.y = 0.05; marker.scale.z = 0.05
        marker.color.a = 1.0; marker.color.r = 1.0; marker.color.g = 0.0; marker.color.b = 0.0
        marker.pose.position.x = float(pos[0]); marker.pose.position.y = float(pos[1]); marker.pose.position.z = float(pos[2])
        self.target_marker_pub.publish(marker)

    def visualize_trajectory(self, points):
        marker = Marker()
        marker.header.frame_id = "base_link"
        marker.type = Marker.SPHERE_LIST; marker.action = Marker.ADD
        marker.scale.x = 0.01; marker.scale.y = 0.01; marker.scale.z = 0.01
        marker.color.a = 1.0; marker.color.r = 1.0; marker.color.g = 1.0
        for p in points:
            pt = Point(); pt.x, pt.y, pt.z = float(p[0]), float(p[1]), float(p[2]); marker.points.append(pt)
        self.traj_pub.publish(marker)

def main(args=None):
    rclpy.init(args=args)
    node = PinnController()
    rclpy.spin(node)
    node.destroy_node()
    rclpy.shutdown()

if __name__ == "__main__":
    main()