#!/usr/bin/env python
import rospy
from std_msgs.msg import Int32
from geometry_msgs.msg import PoseStamped, Pose
from styx_msgs.msg import TrafficLightArray, TrafficLight
from styx_msgs.msg import Lane
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
from light_classification.tl_classifier import TLClassifier
from scipy.spatial import KDTree
import tf
import cv2
import yaml

STATE_COUNT_THRESHOLD = 3

class TLDetector(object):
    def __init__(self):
        rospy.init_node('tl_detector')

        self.pose = None
        self.waypoints = None
        self.waypoint_tree = None
        self.camera_image = None
        self.lights = []

        sub1 = rospy.Subscriber('/current_pose', PoseStamped, self.pose_cb)
        sub2 = rospy.Subscriber('/base_waypoints', Lane, self.waypoints_cb)

        '''
        /vehicle/traffic_lights provides you with the location of the traffic light in 3D map space and
        helps you acquire an accurate ground truth data source for the traffic light
        classifier by sending the current color state of all traffic lights in the
        simulator. When testing on the vehicle, the color state will not be available. You'll need to
        rely on the position of the light and the camera image to predict it.
        '''
        sub3 = rospy.Subscriber('/vehicle/traffic_lights', TrafficLightArray, self.traffic_cb)
        sub6 = rospy.Subscriber('/image_color', Image, self.image_cb)

        config_string = rospy.get_param("/traffic_light_config")
        self.config = yaml.load(config_string)

        self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint', Int32, queue_size=1)

        self.bridge = CvBridge()
        self.light_classifier = TLClassifier()
        self.listener = tf.TransformListener()

        self.state = TrafficLight.UNKNOWN
        self.last_state = TrafficLight.UNKNOWN
        self.last_wp = -1
        self.state_count = 0

        rospy.spin()


    def pose_cb(self, msg):
        self.pose = msg


    def waypoints_cb(self, waypoints):
        self.waypoints = waypoints
        if not self.waypoint_tree:
            self.waypoint_tree = KDTree([
                [w.pose.pose.position.x, w.pose.pose.position.y]
                for w in waypoints.waypoints])


    def traffic_cb(self, msg):
        self.lights = msg.lights


    def image_cb(self, msg):
        """Identifies red lights in the incoming camera image and publishes the index
            of the waypoint closest to the red light's stop line to /traffic_waypoint

        Args:
            msg (Image): image from car-mounted camera
        """
        self.has_image = True
        self.camera_image = msg
        light_wp, state = self.process_traffic_lights()

        # Publish upcoming red lights at camera frequency.
        # Each predicted state has to occur `STATE_COUNT_THRESHOLD` number
        # of times till we start using it. Otherwise the previous stable state is
        # used.

        if self.state != state:
            self.state_count = 0
            self.state = state
        elif self.state_count >= STATE_COUNT_THRESHOLD:
            self.last_state = self.state
            light_wp = light_wp if state == TrafficLight.RED else -1
            self.last_wp = light_wp
            self.upcoming_red_light_pub.publish(Int32(light_wp))
        else:
            self.upcoming_red_light_pub.publish(Int32(self.last_wp))

        if (state == TrafficLight.RED):
            rospy.logwarn("Detected RED light! Count: %i" % self.state_count)
        if (state == TrafficLight.GREEN):
           rospy.logwarn("Detected GREEN light! Count: %i" % self.state_count)
        
        self.state_count += 1


    def get_closest_waypoint(self, x, y):
        """Identifies the closest path waypoint to the given position
            https://en.wikipedia.org/wiki/Closest_pair_of_points_problem

        Args:
            pose (Pose): position to match a waypoint to

        Returns:
            int: index of the closest waypoint in self.waypoints
        """
        return self.waypoint_tree.query([x, y], 1)[1]


    def get_light_state(self, light):
        """Determines the current color of the traffic light

        Args:
            light (TrafficLight): light to classify

        Returns:
            int: ID of traffic light color (specified in styx_msgs/TrafficLight)

        """
        # For testing, just return the light state.
        return light.state

        #if (not self.has_image):
        #    self.prev_light_loc = None
        #    return False

        #cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "bgr8")
        #return self.light_classifier.get_classification(cv_image)


    def process_traffic_lights(self):
        """Finds closest visible traffic light, if one exists, and determines its
            location and color

        Returns:
            int: index of waypoint closest to the upcoming stop line for a traffic light (-1 if none exists)
            int: ID of traffic light color (specified in styx_msgs/TrafficLight)
        """
        closest_light = None
        closest_stop_line_idx = -1 

        stop_line_positions = self.config['stop_line_positions']

        if (self.pose):
            car_waypoint_idx = self.get_closest_waypoint(self.pose.pose.position.x,
                                                         self.pose.pose.position.y)

            # Find the closest visible traffic light.
            # Since the list of lights is short (~8), there isn't much benefit in using KD trees.
            min_distance_to_stop_line = len(self.waypoints.waypoints)
            for light, stop_line in zip(self.lights, stop_line_positions):
                stop_line_idx = self.get_closest_waypoint(stop_line[0], stop_line[1])
                distance_to_stop_line = stop_line_idx - car_waypoint_idx
                # -10 is key to getting the simualtor to work (at least on my not-so-powerful machine).
                # The simulator is laggy and reports the position further ahead that it displays it.
                # Often, when the car is in front of a red traffic light, it doesn't stop because it
                # thinks it's past it.
                if -10 <= distance_to_stop_line <= min_distance_to_stop_line:
                    min_distance_to_stop_line = distance_to_stop_line
                    closest_light = light
                    closest_stop_line_idx = stop_line_idx  

        light_state = self.get_light_state(closest_light) if closest_light else TrafficLight.UNKNOWN
        return closest_stop_line_idx, light_state


if __name__ == '__main__':
    try:
        TLDetector()
    except rospy.ROSInterruptException:
        rospy.logerr('Could not start traffic node.')

#!/usr/bin/env python
import rospy
from std_msgs.msg import Int32
from geometry_msgs.msg import PoseStamped, Pose
from styx_msgs.msg import TrafficLightArray, TrafficLight
from styx_msgs.msg import Lane
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
from light_classification.tl_classifier import TLClassifier
from scipy.spatial import KDTree
import tf
import cv2
import yaml

STATE_COUNT_THRESHOLD = 3

class TLDetector(object):
    def __init__(self):
        rospy.init_node('tl_detector')

        self.pose = None
        self.waypoints = None
        self.waypoint_tree = None
        self.camera_image = None
        self.lights = []

        sub1 = rospy.Subscriber('/current_pose', PoseStamped, self.pose_cb)
        sub2 = rospy.Subscriber('/base_waypoints', Lane, self.waypoints_cb)

        '''
        /vehicle/traffic_lights provides you with the location of the traffic light in 3D map space and
        helps you acquire an accurate ground truth data source for the traffic light
        classifier by sending the current color state of all traffic lights in the
        simulator. When testing on the vehicle, the color state will not be available. You'll need to
        rely on the position of the light and the camera image to predict it.
        '''
        sub3 = rospy.Subscriber('/vehicle/traffic_lights', TrafficLightArray, self.traffic_cb)
        sub6 = rospy.Subscriber('/image_color', Image, self.image_cb)

        config_string = rospy.get_param("/traffic_light_config")
        self.config = yaml.load(config_string)

        self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint', Int32, queue_size=1)

        self.bridge = CvBridge()
        self.light_classifier = TLClassifier()
        self.listener = tf.TransformListener()

        self.state = TrafficLight.UNKNOWN
        self.last_state = TrafficLight.UNKNOWN
        self.last_wp = -1
        self.state_count = 0

        rospy.spin()


    def pose_cb(self, msg):
        self.pose = msg


    def waypoints_cb(self, waypoints):
        self.waypoints = waypoints
        if not self.waypoint_tree:
            self.waypoint_tree = KDTree([
                [w.pose.pose.position.x, w.pose.pose.position.y]
                for w in waypoints.waypoints])


    def traffic_cb(self, msg):
        self.lights = msg.lights


    def image_cb(self, msg):
        """Identifies red lights in the incoming camera image and publishes the index
            of the waypoint closest to the red light's stop line to /traffic_waypoint

        Args:
            msg (Image): image from car-mounted camera
        """
        self.has_image = True
        self.camera_image = msg
        light_wp, state = self.process_traffic_lights()

        # Publish upcoming red lights at camera frequency.
        # Each predicted state has to occur `STATE_COUNT_THRESHOLD` number
        # of times till we start using it. Otherwise the previous stable state is
        # used.

        if self.state != state:
            self.state_count = 0
            self.state = state
        elif self.state_count >= STATE_COUNT_THRESHOLD:
            self.last_state = self.state
            light_wp = light_wp if state == TrafficLight.RED else -1
            self.last_wp = light_wp
            self.upcoming_red_light_pub.publish(Int32(light_wp))
        else:
            self.upcoming_red_light_pub.publish(Int32(self.last_wp))

        if (state == TrafficLight.RED):
            rospy.logwarn("Detected RED light! Count: %i" % self.state_count)
        if (state == TrafficLight.GREEN):
           rospy.logwarn("Detected GREEN light! Count: %i" % self.state_count)
        
        self.state_count += 1


    def get_closest_waypoint(self, x, y):
        """Identifies the closest path waypoint to the given position
            https://en.wikipedia.org/wiki/Closest_pair_of_points_problem

        Args:
            pose (Pose): position to match a waypoint to

        Returns:
            int: index of the closest waypoint in self.waypoints
        """
        return self.waypoint_tree.query([x, y], 1)[1]


    def get_light_state(self, light):
        """Determines the current color of the traffic light

        Args:
            light (TrafficLight): light to classify

        Returns:
            int: ID of traffic light color (specified in styx_msgs/TrafficLight)

        """
        # For testing, just return the light state.
        return light.state

        #if (not self.has_image):
        #    self.prev_light_loc = None
        #    return False

        #cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "bgr8")
        #return self.light_classifier.get_classification(cv_image)


    def process_traffic_lights(self):
        """Finds closest visible traffic light, if one exists, and determines its
            location and color

        Returns:
            int: index of waypoint closest to the upcoming stop line for a traffic light (-1 if none exists)
            int: ID of traffic light color (specified in styx_msgs/TrafficLight)
        """
        closest_light = None
        closest_stop_line_idx = -1 

        stop_line_positions = self.config['stop_line_positions']

        if (self.pose):
            car_waypoint_idx = self.get_closest_waypoint(self.pose.pose.position.x,
                                                         self.pose.pose.position.y)

            # Find the closest visible traffic light.
            # Since the list of lights is short (~8), there isn't much benefit in using KD trees.
            min_distance_to_stop_line = len(self.waypoints.waypoints)
            for light, stop_line in zip(self.lights, stop_line_positions):
                stop_line_idx = self.get_closest_waypoint(stop_line[0], stop_line[1])
                distance_to_stop_line = stop_line_idx - car_waypoint_idx
                # -10 is key to getting the simualtor to work (at least on my not-so-powerful machine).
                # The simulator is laggy and reports the position further ahead that it displays it.
                # Often, when the car is in front of a red traffic light, it doesn't stop because it
                # thinks it's past it.
                if -10 <= distance_to_stop_line <= min_distance_to_stop_line:
                    min_distance_to_stop_line = distance_to_stop_line
                    closest_light = light
                    closest_stop_line_idx = stop_line_idx  

        light_state = self.get_light_state(closest_light) if closest_light else TrafficLight.UNKNOWN
        return closest_stop_line_idx, light_state


if __name__ == '__main__':
    try:
        TLDetector()
    except rospy.ROSInterruptException:
        rospy.logerr('Could not start traffic node.')

#!/usr/bin/env python /vehicle/traffic_lights provides you with the location of the traffic light in 3D map space and         helps you acquire an accurate ground truth data source for the traffic light         classifier by sending the current color state of all traffic lights in the         simulator. When testing on the vehicle, the color state will not be available. You'll need to         rely on the position of the light and the camera image to predict it. Identifies red lights in the incoming camera image and publishes the index             of the waypoint closest to the red light's stop line to /traffic_waypoint          Args:             msg (Image): image from car-mounted camera # Publish upcoming red lights at camera frequency. # Each predicted state has to occur `STATE_COUNT_THRESHOLD` number # of times till we start using it. Otherwise the previous stable state is # used. Identifies the closest path waypoint to the given position             https://en.wikipedia.org/wiki/Closest_pair_of_points_problem          Args:             pose (Pose): position to match a waypoint to          Returns:             int: index of the closest waypoint in self.waypoints Determines the current color of the traffic light          Args:             light (TrafficLight): light to classify          Returns:             int: ID of traffic light color (specified in styx_msgs/TrafficLight) # For testing, just return the light state. #if (not self.has_image): #    self.prev_light_loc = None #    return False #cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "bgr8") #return self.light_classifier.get_classification(cv_image) Finds closest visible traffic light, if one exists, and determines its             location and color          Returns:             int: index of waypoint closest to the upcoming stop line for a traffic light (-1 if none exists)             int: ID of traffic light color (specified in styx_msgs/TrafficLight) # Find the closest visible traffic light. # Since the list of lights is short (~8), there isn't much benefit in using KD trees. # -10 is key to getting the simualtor to work (at least on my not-so-powerful machine). # The simulator is laggy and reports the position further ahead that it displays it. # Often, when the car is in front of a red traffic light, it doesn't stop because it # thinks it's past it.