#! /usr/bin/python
# /home/adam/anaconda3/envs/threatx/bin/python
"""
Node to take in a whole image and perform scene segmentation.
"""

from __future__ import division, print_function
import numpy as np
import cv2

# ROS imports
import rospy
from sensor_msgs.msg import CompressedImage, Image
from cv_bridge import CvBridge
from PIL import Image as PILImage
import StringIO

# Workspace imports
from segmentation.bisenetv2_wrapper import (\
        BiSeNetV2Segmentation,
        real_color_map,
        sim_color_map
        )

bridge = CvBridge()

class BiSeNetV2SegmentationROS(object):
    """
    The ROS implementation of the BiSeNetV2 network. Acts as a segmentation
    network on RGB images.

    """
    def __init__(self, input_image_topic, output_image_topic, bnet):
        self.input_image_topic = input_image_topic
        self.output_image_topic = output_image_topic
        self.bnet = bnet
        self.times = []
        self.frame = 0

    def listener(self):
        """
        The initialization function for this class. When called,
        creates the necessary ROS publishers and subscribers used across
        scopes.

        """
        self.seg_pub = rospy.Publisher(self.output_image_topic, Image,
                queue_size=1)

        if self.input_image_topic[-11:] == '/compressed':
            # If the end of the topic is .../compressed, then it is a
            # compressed image.
            rospy.Subscriber(self.input_image_topic, CompressedImage,
                             lambda msg: self.image_callback_(msg, self.input_image_topic),
                             queue_size=1)
        else:
            rospy.Subscriber(self.input_image_topic, Image,
                             lambda msg: self.image_callback_(msg, self.input_image_topic),
                             queue_size=1)


        self.image_sub = rospy.Subscriber(self.input_image_topic, Image,
                               self.image_callback)
        rospy.loginfo("Waiting for image messages on %s..." % self.input_image_topic)

    def image_callback(self, data):
        """
        The ROS callback to take a ROS Image message and pass through the
        segmentation model. The final output publishes the segmented image.

        :param data: The input image.
        :type data: ROS Image message

        """
        rospy.loginfo("====== Processing image frame %s ======" % self.frame)
        # http://wiki.ros.org/cv_bridge/Tutorials/ConvertingBetweenROSImagesAndOpenCVImagesPython

        # Check whether received message is Image/CompressedImage.
        if hasattr(data, 'format'):
            # Message is CompressedImage.
            sio = StringIO.StringIO(data.data)
            im = PILImage.open(sio)
            img = np.array(im, copy=True)
            img = cv2.cvtColor(img, cv2.COLOR_BAYER_RG2BGR)
        else:
            img = bridge.imgmsg_to_cv2(data, 'rgb8')


        # Trim off alpha channel if it exists
        if img.shape[2] == 4:
            img = img[:, :, :3]
        img = cv2.resize(img, (self.bnet.input_w, self.bnet.input_h))
        tic = rospy.get_time()
        labels = self.bnet([img])
        tic1 = rospy.get_time()
        result_img = np.zeros(img.shape).astype(np.uint8)
        for label, color in sim_color_map.items():
            mask = labels == label
            result_img[mask] = color
        result_img = cv2.addWeighted(img, 0.8, result_img, 0.2, 0)
        result_msg = bridge.cv2_to_imgmsg(result_img, encoding="rgb8")
        self.seg_pub.publish(result_msg)
        self.frame += 1
        self.times.append(tic1 - tic)
        rospy.loginfo("\tSegmentation call averages %.03fs" % np.mean(np.array(self.times)))

def main():
    # Launch the node.
    node = 'segmentation'
    rospy.init_node(node, anonymous=False)
    node_name = rospy.get_name()

    # -------------------------- Read Parameters -----------------------------
    input_image_topic = rospy.get_param('%s/input_image_topic' % node_name, None)
    output_image_topic = rospy.get_param('%s/output_image_topic' % node_name, None)

    b = rospy.get_param('%s/bnet' % node_name, None)

    rospy.loginfo("Constructing segmentation network.")
    bnet = BiSeNetV2Segmentation(b['batch_size'], b['weights_path'],
                                 b['use_trt'], b['input_w'], b['input_h'],
                                 b['num_classes'], None, b['threshold'])
    # ------------------------------------------------------------------------

    BSN = BiSeNetV2SegmentationROS(input_image_topic, output_image_topic, bnet)
    BSN.listener()

    # Blocks until ROS node is shutdown. Yields activity to other threads.
    rospy.spin()

if __name__ == '__main__':
    try:
        main()
    except rospy.ROSInterruptException:
        pass
