#!/usr/bin/env python3

from yolo_detection.msg import BoundingBoxes
from cv_bridge import CvBridge, CvBridgeError
from std_msgs.msg import Bool
from sensor_msgs.msg import Image
from math import floor, ceil
import message_filters
import rospy


class StopSignChecker:
    def __init__(self, stop_time=10, max_sign_distance=5, images_for_clear=10):
        """A check that stops the car a certain time for stop signs

        Makes the car stop on the first image containing a sign, wait a set
        amount of time and then continue driving. The car will then not stop
        again before all stop signs have gone out of sigt for at leas a set
        amount of images after which the system resets - on seeing a new stop
        sign this check will make the car stop again.

        :param stop_time: Duration to remain stopped for after the first
            spotting of a stop sign (seconds)
        :type stop_time: float
        :param max_sign_distance: Signs outside this distance are discarded
            (meters)
        :type max_sign_distance: float
        :param images_for_clear: Once this many images without stop signs have
            been seen in a row the coast is marked clear making way for a new
            observation of a stop sign which would make the car stop again
        :type images_for_clear: int
        """

        self.stop_time = rospy.Duration(stop_time)
        self.max_sign_distance = max_sign_distance
        self.images_for_clear = images_for_clear
        self.sign_first_spotted_time = rospy.Time(0)
        self.images_with_signs = 0
        self.first_sign_of_frame = False

    def new_image(self, time):
        """Callback to be called whenever there is a new camera image"""

        self.images_with_signs = max(self.images_with_signs - 1, 0)
        self.first_sign_of_frame = True

    def process_object(self, distance, obj_class, time):
        """Function to be called for every bounding box in every camera frame"""

        # If it's the first stop sign of the current camera image...
        if (obj_class == 'stop sign' and distance <= self.max_sign_distance and
                self.first_sign_of_frame):
            self.first_sign_of_frame = False

            # ... and it's the first stop sign we se after a period of no stop
            # sign...
            if self.images_with_signs == 0:
                # ... Mark this as the time when we first spotted a sign
                self.sign_first_spotted_time = time

            # Increase by two here since we also decrease by one every frame in
            # new_image(): 2 + (-1) = 1, meaning we increment images_with_signs
            # by one every frame where there is atleast one stop sign
            self.images_with_signs = min(self.images_with_signs + 2,
                                         self.images_for_clear)

    def should_stop(self, time):
        """Determines if the car should stop according to this check"""

        return time <= self.sign_first_spotted_time + self.stop_time


class FullStopChecker:
    def __init__(self, classes, remain_stopped_time=2, max_obj_distance=3.2):
        """A check that completely stops the car for certain objects

        :param classes: List of classes to stop for
        :type classes: list
        :param remain_stopped_time: Duration of time to remain stopped for
            after all objects have gone out of range (seconds)
        :type remain_stopped_time: float
        :param max_obj_distance: Objects outside this distance are discarded
            (meters)
        :type max_obj_distance: float
        """

        self.classes = classes
        self.remain_stopped_time = rospy.Duration(remain_stopped_time)
        self.max_obj_distance = max_obj_distance
        self.obj_last_spotted_time = rospy.Time(0)

    def new_image(self, time):
        """Callback to be called whenever there is a new camera image"""

        pass

    def process_object(self, distance, obj_class, time):
        """Function to be called for every bounding box in every camera frame"""

        if distance <= self.max_obj_distance and obj_class in self.classes:
            self.obj_last_spotted_time = time

    def should_stop(self, time):
        """Determines if the car should stop according to this check"""

        return time <= self.obj_last_spotted_time + self.remain_stopped_time


def main():
    # Initialize ROS node
    rospy.init_node('object_avoidance', anonymous=True)

    # Create publisher for when there is an object in front of the camera
    pub = rospy.Publisher('object_in_front', Bool, queue_size=1)

    # CvBridge is used to convert the ROS image to numpy arrays
    bridge = CvBridge()

    # List of different checks to perform
    checks = [FullStopChecker(['person', 'dog', 'cat', 'horse']),
              StopSignChecker()]

    def should_stop(time=None):
        """Whether the checks are telling the car to stop

        :param time: ROS Time to perform the check at or None to check for the
            current time
        """

        time = rospy.Time.now() if time is None else time
        for check in checks:
            if check.should_stop(time):
                return True
        return False

    def process_image(img, boxes):
        time = rospy.Time.now()

        try:
            # Convert from ROS Image message to numpy array
            cv_image = bridge.imgmsg_to_cv2(img, desired_encoding='16UC1')
        except CvBridgeError as e:
            rospy.logwarn(e)
            return

        # Signal to all checks that a new camera image has arrived
        for check in checks:
            check.new_image(time)

        for box in boxes.bounding_boxes:
            # Crop the bounding boxes a little further to increase distance
            # measurement accuracy
            xmargin = int((box.xmax - box.xmin) * 0.25)
            ymargin = int((box.ymax - box.ymin) * 0.25)
            xmin = box.xmin + xmargin
            xmax = box.xmax - xmargin
            ymin = box.ymin + ymargin
            ymax = box.ymax - ymargin

            # We can only calculate a distance to an object if it takes up
            # space in the depth image
            if xmin == xmax or ymin == ymax:
                continue

            # Convert from millimeters to meters
            distance = cv_image[ymin:ymax, xmin:xmax].mean() / 1000.0

            # Perform per-bonding-box logic for every type of check
            for check in checks:
                check.process_object(distance, box.Class, time)

        # Publish potential stop signal immediately to avoid delay
        pub.publish(should_stop(time))

    # Subscribe to the image given by the realsense camera and the bounding
    # boxes provided by YOLO then align the messages by their timestamps
    image_sub = message_filters.Subscriber('image', Image, buff_size=2**24)
    boxes_sub = message_filters.Subscriber('boxes', BoundingBoxes)
    ts = message_filters.TimeSynchronizer([image_sub, boxes_sub], 10)
    ts.registerCallback(process_image)

    rate = rospy.Rate(1)
    while not rospy.is_shutdown():
        # Publish stop (or go) signal periodically to the system does not
        # freeze when we stop recieving frames from YOLO (when the camera image
        # does not contain any objects recognized by YOLO)
        pub.publish(should_stop())

        rate.sleep()


if __name__ == '__main__':
    main()
