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")
        # 注意：这里 output_dim 必须是 9 (三角编码版)
        self.model = PINN_IK(output_dim=9).to(self.device)
        self.stats = None
        
        # 确保这里的文件名和你 train.py 最后保存的文件名一致
        # 上一次如果是 best_model_smooth.pth 就用这个
        model_path = os.path.join(current_dir, "best_model_smooth.pth")
        
        if os.path.exists(model_path):
            # =======================================================
            # 🔧 修复点：添加 weights_only=False
            # =======================================================
            checkpoint = torch.load(model_path, map_location=self.device, weights_only=False)
            
            self.model.load_state_dict(checkpoint['model_state'])
            self.stats = checkpoint['stats']
            self.model.eval()
            self.get_logger().info(f"✅ 模型加载成功: {model_path}")
        else:
            self.get_logger().error(f"❌ 找不到模型: {model_path}")

        # 验证用的 FK 工具
        self.fk_solver = SoftArmKinematics(3, 0.24, 3, 0.033)

        # 2. 通信接口
        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("等待目标指令... (请在 Rviz 中点击)")

    def decode_output(self, raw_output):
        """
        解码: [t1, cos1, sin1, ...] -> [t1, phi1, ...]
        """
        # 反归一化
        y_mean = self.stats['y_mean']
        y_std = self.stats['y_std']
        real_vals = raw_output * y_std + y_mean
        
        final_angles = []
        for i in range(3):
            theta = real_vals[i*3 + 0]
            cos_v = real_vals[i*3 + 1]
            sin_v = real_vals[i*3 + 2]
            
            # 三角解码 phi = atan2(sin, cos)
            phi = np.arctan2(sin_v, cos_v)
            
            final_angles.append(theta)
            final_angles.append(phi)
            
        return np.array(final_angles)

    def goal_callback(self, msg):
        # 强制 Z=0.5 进行测试，防止点到地面
        target_pos = np.array([msg.pose.position.x, msg.pose.position.y, 0.5])
        
        dist = np.linalg.norm(target_pos)
        if dist > 0.75:
            self.get_logger().warn(f"目标太远 ({dist:.2f}m)，忽略")
            return

        self.get_logger().info(f"目标: {target_pos}，开始规划")
        self.visualize_target(target_pos)
        threading.Thread(target=self.execute_move, args=(target_pos,)).start()

    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_decoded = None
        
        for pt in trajectory:
            with torch.no_grad():
                # 预处理输入 (归一化)
                X_mean = self.stats['X_mean']; X_std = self.stats['X_std']
                norm_input = (pt - X_mean) / X_std
                tensor_input = torch.FloatTensor(norm_input).to(self.device)
                
                # 推理
                norm_output = self.model(tensor_input).numpy()
                
                # 解码输出
                real_angles = self.decode_output(norm_output)
                final_angles_decoded = real_angles
            
            # 发送指令
            msg = Float64MultiArray()
            msg.data = real_angles.tolist()
            self.cmd_pub.publish(msg)
            time.sleep(0.04) # 控制速度
            
        self.current_pos = target_pos
        
        # 验证精度
        if final_angles_decoded is not None:
            self.check_accuracy(target_pos, final_angles_decoded)

    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(f"\n>>> 精度验证 <<<\n目标: {target_pos}\n实际: {actual_pos}\n误差: {error:.4f} m (Loss较低时此误差应很小)\n")

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