#!/usr/bin/env python3

import rospy
import random
from geometry_msgs.msg import PointStamped, Point
from visualization_msgs.msg import Marker, MarkerArray
from Service.srv import FFT

MAX_POINTS = 100
FFT_WINDOW_SIZE = 128

original_trajectory = []
noisy_trajectory = []
noisy_x_history = []
noisy_y_history = []
current_x_magnitudes = []
current_y_magnitudes = []
noisy_pub = None
marker_pub = None
fft_client = None

def position_callback(msg):
    global original_trajectory, noisy_trajectory
    global noisy_x_history, noisy_y_history
    global current_x_magnitudes, current_y_magnitudes
    global noisy_pub, marker_pub, fft_client

    noise_x = random.uniform(-0.1, 0.1)    # 产生 -0.1 ～0.1 的随机噪声
    noise_y = random.uniform(-0.1, 0.1)

    noisy_msg = PointStamped()
    noisy_msg.header = msg.header
    noisy_msg.point.x = msg.point.x + noise_x    # 在原信号上叠加噪声
    noisy_msg.point.y = msg.point.y + noise_y
    noisy_msg.point.z = 0.0

    original_trajectory.append(msg.point)    # 原点轨迹
    if len(original_trajectory) > MAX_POINTS:
        original_trajectory.pop(0)
    noisy_trajectory.append(noisy_msg.point)    #加噪轨迹
    if len(noisy_trajectory) > MAX_POINTS:
        noisy_trajectory.pop(0)

    noisy_x_history.append(noisy_msg.point.x)    #保存历史坐标信息
    noisy_y_history.append(noisy_msg.point.y)

    marker_array = MarkerArray()

    # 原始信号轨迹绘制
    original_line = Marker()
    original_line.header = msg.header
    original_line.ns = "trajectory"
    original_line.id = 0
    original_line.type = Marker.LINE_STRIP
    original_line.action = Marker.ADD
    original_line.scale.x = 0.02
    original_line.color.g = 1.0    # 绿色线
    original_line.color.a = 1.0
    original_line.points = list(original_trajectory)
    marker_array.markers.append(original_line)

    # 加噪信号轨迹绘制
    noisy_line = Marker()
    noisy_line.header = msg.header
    noisy_line.ns = "trajectory"
    noisy_line.id = 1
    noisy_line.type = Marker.LINE_STRIP
    noisy_line.action = Marker.ADD
    noisy_line.scale.x = 0.02
    noisy_line.color.r = 1.0    # 红色线
    noisy_line.color.a = 1.0
    noisy_line.points = list(noisy_trajectory)
    marker_array.markers.append(noisy_line)

    # 原始信号位置点
    original_point = Marker()
    original_point.header = msg.header
    original_point.ns = "current_point"
    original_point.id = 0
    original_point.type = Marker.SPHERE
    original_point.action = Marker.ADD
    original_point.pose.position = msg.point
    original_point.pose.orientation.w = 1.0
    original_point.scale.x = 0.1
    original_point.scale.y = 0.1
    original_point.scale.z = 0.1
    original_point.color.g = 1.0    # 绿色球
    original_point.color.a = 1.0
    marker_array.markers.append(original_point)

    # 加噪信号位置点
    noisy_point = Marker()
    noisy_point.header = msg.header
    noisy_point.ns = "current_point"
    noisy_point.id = 1
    noisy_point.type = Marker.SPHERE
    noisy_point.action = Marker.ADD
    noisy_point.pose.position = noisy_msg.point
    noisy_point.pose.orientation.w = 1.0
    noisy_point.scale.x = 0.1
    noisy_point.scale.y = 0.1
    noisy_point.scale.z = 0.1
    noisy_point.color.r = 1.0    # 红色球
    noisy_point.color.a = 1.0
    marker_array.markers.append(noisy_point)

    # 当点数 >=128 时利用快速傅里叶变换FFT计算加噪信号频谱
    if len(noisy_x_history) >= FFT_WINDOW_SIZE:
        try:
            res = fft_client(data_x=noisy_x_history[-FFT_WINDOW_SIZE:],
                             data_y=noisy_y_history[-FFT_WINDOW_SIZE:],
                             window_size=FFT_WINDOW_SIZE)
            rospy.loginfo("FFT计算成功")
            current_x_magnitudes[:] = res.magnitude_x
            current_y_magnitudes[:] = res.magnitude_y
        except rospy.ServiceException as e:
            rospy.logerr("FFT计算失败: %s" % e)
    
    # 判断频谱，生成频谱图
    if current_x_magnitudes and current_y_magnitudes:
        x_spectrum = Marker()
        x_spectrum.header = msg.header
        x_spectrum.ns = "x_spectrum"
        x_spectrum.id = 2
        x_spectrum.type = Marker.CUBE_LIST
        x_spectrum.action = Marker.ADD
        x_spectrum.scale.x = 0.08
        x_spectrum.scale.y = 0.01
        x_spectrum.scale.z = 1.0
        x_spectrum.color.b = 1.0
        x_spectrum.color.a = 0.8

        y_spectrum = Marker()
        y_spectrum.header = msg.header
        y_spectrum.ns = "y_spectrum"
        y_spectrum.id = 3
        y_spectrum.type = Marker.CUBE_LIST
        y_spectrum.action = Marker.ADD
        y_spectrum.scale.x = 0.08
        y_spectrum.scale.y = 0.01
        y_spectrum.scale.z = 1.0
        y_spectrum.color.r = 1.0
        y_spectrum.color.a = 0.8

        for i in range(len(current_x_magnitudes)):
            px = Point()
            px.x = i * 0.1
            px.y = 2.0
            px.z = (current_x_magnitudes[i] / float(FFT_WINDOW_SIZE)) * 10.0
            x_spectrum.points.append(px)
            py = Point()
            py.x = px.x
            py.y = 1.0
            py.z = (current_y_magnitudes[i] / float(FFT_WINDOW_SIZE)) * 10.0
            y_spectrum.points.append(py)

        marker_array.markers.append(x_spectrum)
        marker_array.markers.append(y_spectrum)

        text_x = Marker()
        text_x.header = msg.header
        text_x.ns = "spectrum_labels"
        text_x.type = Marker.TEXT_VIEW_FACING
        text_x.action = Marker.ADD
        text_x.scale.z = 0.3
        text_x.color.r = 1.0
        text_x.color.g = 1.0
        text_x.color.b = 1.0
        text_x.color.a = 1.0
        text_x.id = 10
        text_x.pose.position.x = -0.5
        text_x.pose.position.y = 2.0
        text_x.pose.position.z = 1.0
        text_x.text = "坐标X频谱"
        marker_array.markers.append(text_x)
        
        text_y = Marker()
        text_y.header = msg.header
        text_y.ns = "spectrum_labels"
        text_y.type = Marker.TEXT_VIEW_FACING
        text_y.action = Marker.ADD
        text_y.scale.z = 0.3
        text_y.color.r = 1.0
        text_y.color.g = 1.0
        text_y.color.b = 1.0
        text_y.color.a = 1.0
        text_y.id = 11
        text_y.pose.position.x = -0.5
        text_y.pose.position.y = 1.0
        text_y.pose.position.z = 1.0
        text_y.text = "坐标Y频谱"
        marker_array.markers.append(text_y)

    noisy_pub.publish(noisy_msg)
    marker_pub.publish(marker_array)

    rospy.loginfo("订阅端接收坐标信息: x=%.3f, y=%.3f | 加噪坐标信息: x=%.3f, y=%.3f",
                  msg.point.x, msg.point.y, noisy_msg.point.x, noisy_msg.point.y)

if __name__ == '__main__':
    rospy.init_node('Subscriber_Client_S')
    rospy.Subscriber('position', PointStamped, position_callback)

    noisy_pub = rospy.Publisher('noisy_position', PointStamped, queue_size=10)
    marker_pub = rospy.Publisher('visualization_marker_array', MarkerArray, queue_size=10)
    fft_client = rospy.ServiceProxy('fft_service', FFT)
    
    rospy.loginfo("等待快速傅里叶FFT服务端启动")
    rospy.wait_for_service('fft_service')
    rospy.loginfo("FFT服务连接成功")
    rospy.spin()