From ec7b3b63fda968aa79ef7b75f340f52a4ec1cc72 Mon Sep 17 00:00:00 2001
From: Axel <axellopez92@outlook.com>
Date: Wed, 16 Sep 2020 13:09:18 +0200
Subject: [PATCH] Added off state to traffic lights.

---
 .../Carla/Traffic/TrafficLightComponent.cpp     | 17 +++++++++++++++--
 .../Source/Carla/Traffic/TrafficLightState.h    |  8 +++++++-
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/Unreal/CarlaUE4/Plugins/Carla/Source/Carla/Traffic/TrafficLightComponent.cpp b/Unreal/CarlaUE4/Plugins/Carla/Source/Carla/Traffic/TrafficLightComponent.cpp
index baaea9ddc..9b0502875 100644
--- a/Unreal/CarlaUE4/Plugins/Carla/Source/Carla/Traffic/TrafficLightComponent.cpp
+++ b/Unreal/CarlaUE4/Plugins/Carla/Source/Carla/Traffic/TrafficLightComponent.cpp
@@ -97,9 +97,16 @@ void UTrafficLightComponent::SetLightState(ETrafficLightState NewState)
       {
         Controller->SetTrafficLight(nullptr);
       }
+      // workarround for tm not supporting off state
+      if (LightState == ETrafficLightState::Off)
+      {
+        Controller->SetTrafficLightState(ETrafficLightState::Green);
+        Controller->SetTrafficLight(nullptr);
+      }
     }
   }
-  if (LightState == ETrafficLightState::Green)
+  if (LightState == ETrafficLightState::Green ||
+      LightState == ETrafficLightState::Off)
   {
     Vehicles.Empty();
   }
@@ -153,12 +160,18 @@ void UTrafficLightComponent::OnOverlapTriggerBox(UPrimitiveComponent *Overlapped
     if (VehicleController)
     {
       VehicleController->SetTrafficLightState(LightState);
-      if (LightState != ETrafficLightState::Green)
+      if (LightState != ETrafficLightState::Green &&
+          LightState != ETrafficLightState::Off)
       {
         Vehicles.Add(VehicleController);
         VehicleController->SetTrafficLight(
             Cast<ATrafficLightBase>(GetOwner()));
       }
+      // workarround for tm not supporting off state
+      if (LightState == ETrafficLightState::Off)
+      {
+        VehicleController->SetTrafficLightState(ETrafficLightState::Green);
+      }
     }
   }
 }
diff --git a/Unreal/CarlaUE4/Plugins/Carla/Source/Carla/Traffic/TrafficLightState.h b/Unreal/CarlaUE4/Plugins/Carla/Source/Carla/Traffic/TrafficLightState.h
index daba3d2e7..cc5b26fd7 100644
--- a/Unreal/CarlaUE4/Plugins/Carla/Source/Carla/Traffic/TrafficLightState.h
+++ b/Unreal/CarlaUE4/Plugins/Carla/Source/Carla/Traffic/TrafficLightState.h
@@ -6,11 +6,17 @@
 
 #pragma once
 
+#include <compiler/disable-ue4-macros.h>
+#include <carla/rpc/TrafficLightState.h>
+#include <compiler/enable-ue4-macros.h>
+
 #include "TrafficLightState.generated.h"
 
+
 UENUM(BlueprintType)
 enum class ETrafficLightState : uint8 {
   Red       UMETA(DisplayName = "Red"),
   Yellow    UMETA(DisplayName = "Yellow"),
-  Green     UMETA(DisplayName = "Green")
+  Green     UMETA(DisplayName = "Green"),
+  Off       UMETA(DisplayName = "Off")
 };