From 144c20d58375c123a747ec26a57d43c039be60fc Mon Sep 17 00:00:00 2001
From: KhalilSelyan <khalilselyan2@gmail.com>
Date: Mon, 15 Jan 2024 11:53:04 +0300
Subject: [PATCH] Update traffic light display and data handling with new
 message type to avoid Qt MOC collision `signals` keyword

---
 .../include/traffic_display.hpp               |  7 ++-
 .../src/traffic_display.cpp                   | 49 +++++++++++++------
 2 files changed, 38 insertions(+), 18 deletions(-)

diff --git a/common/awf_vehicle_rviz_plugin/awf_2d_overlay_vehicle/include/traffic_display.hpp b/common/awf_vehicle_rviz_plugin/awf_2d_overlay_vehicle/include/traffic_display.hpp
index 48ad17f7d5061..46524aa2614f5 100644
--- a/common/awf_vehicle_rviz_plugin/awf_2d_overlay_vehicle/include/traffic_display.hpp
+++ b/common/awf_vehicle_rviz_plugin/awf_2d_overlay_vehicle/include/traffic_display.hpp
@@ -25,12 +25,11 @@ class TrafficDisplay
 public:
   TrafficDisplay();
   void drawTrafficLightIndicator(QPainter & painter, const QRectF & backgroundRect);
-  // void updateTrafficLightData(const
-  // autoware_perception_msgs::msg::TrafficSignalArray::ConstSharedPtr msg);
-  rviz_2d_overlay_msgs::msg::TrafficSignalArrayUI temp;
+  void updateTrafficLightData(
+    const rviz_2d_overlay_msgs::msg::TrafficSignalArrayUI::ConstSharedPtr & msg);
+  rviz_2d_overlay_msgs::msg::TrafficSignalArrayUI current_traffic_;
 
 private:
-  int current_traffic_;  // Internal variable to store current gear
   QImage traffic_light_image_;
   // yellow #CFC353
   QColor yellow = QColor(207, 195, 83);
diff --git a/common/awf_vehicle_rviz_plugin/awf_2d_overlay_vehicle/src/traffic_display.cpp b/common/awf_vehicle_rviz_plugin/awf_2d_overlay_vehicle/src/traffic_display.cpp
index 6919b6a74be0f..32ab6c61c9b95 100644
--- a/common/awf_vehicle_rviz_plugin/awf_2d_overlay_vehicle/src/traffic_display.cpp
+++ b/common/awf_vehicle_rviz_plugin/awf_2d_overlay_vehicle/src/traffic_display.cpp
@@ -19,20 +19,16 @@
 namespace awf_2d_overlay_vehicle
 {
 
-TrafficDisplay::TrafficDisplay() : current_traffic_(0)
+TrafficDisplay::TrafficDisplay()
 {
   traffic_light_image_.load(":/assets/images/traffic.png");
 }
 
-// void TrafficDisplay::updateTrafficLightData(const
-// autoware_perception_msgs::msg::TrafficSignalArray::ConstSharedPtr msg)
-// {
-//     // Update the internal variable
-//     // current_traffic_ = msg->traffic_signal_id;
-
-//     // RCLCPP_INFO(rclcpp::get_logger("rcl"), "Traffic Light: %d", msg->elements.size());
-//     // Force the plugin to redraw
-// }
+void TrafficDisplay::updateTrafficLightData(
+  const rviz_2d_overlay_msgs::msg::TrafficSignalArrayUI::ConstSharedPtr & msg)
+{
+  current_traffic_ = *msg;
+}
 
 void TrafficDisplay::drawTrafficLightIndicator(QPainter & painter, const QRectF & backgroundRect)
 {
@@ -53,12 +49,37 @@ void TrafficDisplay::drawTrafficLightIndicator(QPainter & painter, const QRectF
   QRectF imageRect =
     circleRect.adjusted(15, 15, -15, -15);  // Adjusting the rectangle to make the image smaller
 
+  QImage scaled_traffic_image = traffic_light_image_.scaled(
+    imageRect.size().toSize(), Qt::KeepAspectRatio, Qt::SmoothTransformation);
+
+  if (current_traffic_.traffic_signals.size() > 0) {
+    switch (current_traffic_.traffic_signals[0].elements[0].color) {
+      case 1:
+        painter.setBrush(QBrush(red));
+        painter.drawEllipse(circleRect.center(), 30, 30);
+        break;
+      case 2:
+        painter.setBrush(QBrush(yellow));
+        painter.drawEllipse(circleRect.center(), 30, 30);
+        break;
+      case 3:
+        painter.setBrush(QBrush(green));
+        painter.drawEllipse(circleRect.center(), 30, 30);
+        break;
+      case 4:
+        painter.setBrush(QBrush(gray));
+        painter.drawEllipse(circleRect.center(), 30, 30);
+        break;
+      default:
+        painter.setBrush(QBrush(gray));
+        painter.drawEllipse(circleRect.center(), 30, 30);
+        break;
+    }
+  }
   // make the image thicker
   painter.setPen(QPen(Qt::black, 2, Qt::SolidLine, Qt::RoundCap, Qt::RoundJoin));
-  // Draw the traffic light image on top of the circle
-  painter.drawImage(
-    imageRect, traffic_light_image_.scaled(
-                 imageRect.size().toSize(), Qt::KeepAspectRatio, Qt::SmoothTransformation));
+
+  painter.drawImage(imageRect, scaled_traffic_image);
 }
 
 QImage TrafficDisplay::coloredImage(const QImage & source, const QColor & color)