#include <ros/ros.h>
#include <tf/transform_broadcaster.h>
#include <nav_msgs/Odometry.h>

double vx = 0.0;
double vy = 0.0;
double vth = 0.0;

void odomCallback(const nav_msgs::Odometry::ConstPtr& msg) {
  // Extract velocity data from the /odom topic
  vx = msg->twist.twist.linear.x;
  vy = msg->twist.twist.linear.y;
  vth = msg->twist.twist.angular.z;
}

int main(int argc, char** argv) {
  ros::init(argc, argv, "odometry_publisher");

  ros::NodeHandle n;

  // Subscriber for /odom topic
  ros::Subscriber odom_sub = n.subscribe("odom", 50, odomCallback);

  // Publisher for odom topic
  ros::Publisher odom_pub = n.advertise<nav_msgs::Odometry>("odom_out", 50);

  tf::TransformBroadcaster odom_broadcaster;

  double x = 0.0;
  double y = 0.0;
  double th = 0.0;

  ros::Time current_time, last_time;
  current_time = ros::Time::now();
  last_time = ros::Time::now();

  ros::Rate r(50.0); // Update rate increased to 50 Hz

  while (ros::ok()) {
    ros::spinOnce();  // Check for incoming messages
    current_time = ros::Time::now();

    // Compute odometry based on subscribed velocity data
    double dt = (current_time - last_time).toSec();
    double delta_x = (vx * cos(th)) * dt;
    double delta_y = (vx * sin(th)) * dt;
    double delta_th = vth * dt;

    x += delta_x;
    y += delta_y;
    th += delta_th;

    // Create quaternion from yaw
    geometry_msgs::Quaternion odom_quat = tf::createQuaternionMsgFromYaw(th);

    // Publish the transform over tf
    geometry_msgs::TransformStamped odom_trans;
    odom_trans.header.stamp = current_time;
    odom_trans.header.frame_id = "odom";
    odom_trans.child_frame_id = "base_link";

    odom_trans.transform.translation.x = x;
    odom_trans.transform.translation.y = y;
    odom_trans.transform.translation.z = 0.0;
    odom_trans.transform.rotation = odom_quat;

    odom_broadcaster.sendTransform(odom_trans);

    // Publish the odometry message
    nav_msgs::Odometry odom;
    odom.header.stamp = current_time;
    odom.header.frame_id = "odom";

    // Set the position
    odom.pose.pose.position.x = x;
    odom.pose.pose.position.y = y;
    odom.pose.pose.position.z = 0.0;
    odom.pose.pose.orientation = odom_quat;

    // Set the velocity
    odom.child_frame_id = "base_link";
    odom.twist.twist.linear.x = vx;
    odom.twist.twist.linear.y = vy;
    odom.twist.twist.angular.z = vth;

    odom_pub.publish(odom);

    last_time = current_time;
    r.sleep();
  }
}