#include <ros/ros.h>
#include <geometry_msgs/PointStamped.h>
#include <visualization_msgs/Marker.h>
#include <visualization_msgs/MarkerArray.h>
#include <shiyan2/FFT.h>
#include <std_msgs/Float64.h>
#include <random>
#include <vector>

// 全局变量
ros::Publisher noisy_pub;
ros::Publisher marker_pub;
ros::Publisher spectrum_pub;  // 新增频谱可视化发布器
ros::ServiceClient fft_client;

// 随机数生成器
std::default_random_engine generator;
std::uniform_real_distribution<double> distribution(-0.1, 0.1);

// 轨迹存储
std::vector<geometry_msgs::Point> original_trajectory;
std::vector<geometry_msgs::Point> noisy_trajectory;
std::vector<double> noisy_x_history;
std::vector<double> noisy_y_history;
const int MAX_POINTS = 100;
const int FFT_WINDOW_SIZE = 128;

// 新增：频谱数据存储
std::vector<double> current_x_magnitudes;
std::vector<double> current_y_magnitudes;

void positionCallback(const geometry_msgs::PointStamped::ConstPtr& msg) {
    // 1. 生成噪声数据
    double noise_x = distribution(generator);
    double noise_y = distribution(generator);
    
    // 2. 创建带噪声的点消息
    geometry_msgs::PointStamped noisy_msg;
    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;
    
    // 3. 更新轨迹
    original_trajectory.push_back(msg->point);
    if (original_trajectory.size() > MAX_POINTS) {
        original_trajectory.erase(original_trajectory.begin());
    }
    
    noisy_trajectory.push_back(noisy_msg.point);
    if (noisy_trajectory.size() > MAX_POINTS) {
        noisy_trajectory.erase(noisy_trajectory.begin());
    }
    
    // 4. 更新FFT数据缓冲区
    noisy_x_history.push_back(noisy_msg.point.x);
    noisy_y_history.push_back(noisy_msg.point.y);
    
    // 5. 创建可视化标记数组
    visualization_msgs::MarkerArray marker_array;
    
    // 5.1 原始轨迹 (绿色线条)
    visualization_msgs::Marker original_line;
    original_line.header = msg->header;
    original_line.ns = "trajectory";
    original_line.id = 0;
    original_line.type = visualization_msgs::Marker::LINE_STRIP;
    original_line.action = visualization_msgs::Marker::ADD;
    original_line.scale.x = 0.02;
    original_line.color.g = 1.0;
    original_line.color.a = 1.0;
    original_line.points = original_trajectory;
    marker_array.markers.push_back(original_line);
    
    // 5.2 噪声轨迹 (红色线条)
    visualization_msgs::Marker noisy_line;
    noisy_line.header = msg->header;
    noisy_line.ns = "trajectory";
    noisy_line.id = 1;
    noisy_line.type = visualization_msgs::Marker::LINE_STRIP;
    noisy_line.action = visualization_msgs::Marker::ADD;
    noisy_line.scale.x = 0.02;
    noisy_line.color.r = 1.0;
    noisy_line.color.a = 1.0;
    noisy_line.points = noisy_trajectory;
    marker_array.markers.push_back(noisy_line);
    
    // 5.3 当前原始点 (绿色球体)
    visualization_msgs::Marker original_point;
    original_point.header = msg->header;
    original_point.ns = "current_point";
    original_point.id = 0;
    original_point.type = visualization_msgs::Marker::SPHERE;
    original_point.action = visualization_msgs::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.push_back(original_point);
    
    // 5.4 当前噪声点 (红色球体)
    visualization_msgs::Marker noisy_point;
    noisy_point.header = msg->header;
    noisy_point.ns = "current_point";
    noisy_point.id = 1;
    noisy_point.type = visualization_msgs::Marker::SPHERE;
    noisy_point.action = visualization_msgs::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.push_back(noisy_point);
    
    // 6. 当数据达到FFT窗口大小时计算频谱
    if (noisy_x_history.size() >= FFT_WINDOW_SIZE) {
        // 准备FFT请求
        shiyan2::FFT fft_srv;
        fft_srv.request.data_x.assign(noisy_x_history.end() - FFT_WINDOW_SIZE, noisy_x_history.end());
        fft_srv.request.data_y.assign(noisy_y_history.end() - FFT_WINDOW_SIZE, noisy_y_history.end());
        fft_srv.request.window_size = FFT_WINDOW_SIZE;
        
        // 调用FFT服务
        if (fft_client.call(fft_srv)) {
            ROS_INFO("FFT computed successfully");
            
            // 保存频谱结果用于可视化
            current_x_magnitudes = fft_srv.response.magnitude_x;
            current_y_magnitudes = fft_srv.response.magnitude_y;
        } else {
            ROS_ERROR("Failed to call FFT service");
        }
    }
    
    // 7. 创建频谱可视化
    if (!current_x_magnitudes.empty() && !current_y_magnitudes.empty()) {
        // 7.1 X方向频谱柱状图 (蓝色)
        visualization_msgs::Marker x_spectrum;
        x_spectrum.header = msg->header;
        x_spectrum.ns = "x_spectrum";
        x_spectrum.id = 2;  // 唯一ID
        x_spectrum.type = visualization_msgs::Marker::CUBE_LIST;
        x_spectrum.action = visualization_msgs::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;   // 透明度
        
        // 7.2 Y方向频谱柱状图 (红色)
        visualization_msgs::Marker y_spectrum;
        y_spectrum.header = msg->header;
        y_spectrum.ns = "y_spectrum";
        y_spectrum.id = 3;  // 唯一ID
        y_spectrum.type = visualization_msgs::Marker::CUBE_LIST;
        y_spectrum.action = visualization_msgs::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;
        
        // 7.3 设置柱子位置 (X频谱放在Y=2.0平面，Y频谱放在Y=1.0平面)
        for (int i = 0; i < current_x_magnitudes.size(); ++i) {
            geometry_msgs::Point p;
            p.x = i * 0.1;                   // X位置 (频率)
            p.y = 2.0;                       // Y位置 (固定平面)
            p.z = current_x_magnitudes[i] / FFT_WINDOW_SIZE * 10;  // 缩放幅度值
            x_spectrum.points.push_back(p);
            
            p.y = 1.0;                       // Y位置 (固定平面)
            p.z = current_y_magnitudes[i] / FFT_WINDOW_SIZE * 10;  // 缩放幅度值
            y_spectrum.points.push_back(p);
        }
        
        // 7.4 将频谱标记加入数组
        marker_array.markers.push_back(x_spectrum);
        marker_array.markers.push_back(y_spectrum);
        
        // 7.5 添加坐标轴说明
        visualization_msgs::Marker text_marker;
        text_marker.header = msg->header;
        text_marker.ns = "spectrum_labels";
        text_marker.type = visualization_msgs::Marker::TEXT_VIEW_FACING;
        text_marker.action = visualization_msgs::Marker::ADD;
        text_marker.scale.z = 0.3;
        text_marker.color.r = 1.0;
        text_marker.color.g = 1.0;
        text_marker.color.b = 1.0;
        text_marker.color.a = 1.0;
        
        // X频谱标签
        text_marker.id = 10;
        text_marker.pose.position.x = -0.5;
        text_marker.pose.position.y = 2.0;
        text_marker.pose.position.z = 1.0;
        text_marker.text = "X Spectrum";
        marker_array.markers.push_back(text_marker);
        
        // Y频谱标签
        text_marker.id = 11;
        text_marker.pose.position.x = -0.5;
        text_marker.pose.position.y = 1.0;
        text_marker.pose.position.z = 1.0;
        text_marker.text = "Y Spectrum";
        marker_array.markers.push_back(text_marker);
    }
    
    // 8. 发布消息
    noisy_pub.publish(noisy_msg);
    marker_pub.publish(marker_array);
    
    // 9. 打印信息
    ROS_INFO("Published: x=%.3f, y=%.3f | Noisy: x=%.3f, y=%.3f",
             msg->point.x, msg->point.y, 
             noisy_msg.point.x, noisy_msg.point.y);
}

int main(int argc, char** argv) {
    ros::init(argc, argv, "nodeB");
    ros::NodeHandle nh;
    
    // 订阅原始位置
    ros::Subscriber sub = nh.subscribe("position", 10, positionCallback);
    
    // 发布带噪声的位置
    noisy_pub = nh.advertise<geometry_msgs::PointStamped>("noisy_position", 10);
    
    // 发布可视化标记数组
    marker_pub = nh.advertise<visualization_msgs::MarkerArray>("visualization_marker_array", 10);
    
    // 创建FFT服务客户端
    fft_client = nh.serviceClient<shiyan2::FFT>("fft_service");
    
    // 等待服务可用
    if (ros::service::waitForService("fft_service", ros::Duration(5.0))) {
        ROS_INFO("FFT service connected");
    } else {
        ROS_WARN("Waiting for FFT service...");
        ros::service::waitForService("fft_service");
        ROS_INFO("FFT service connected");
    }
    
    ros::spin();
    return 0;
}