#! /usr/bin/python
"""
node for viewing sections of imagery
"""
from __future__ import division, print_function
import numpy as np
import threading
from collections import deque
import cv2
import time
from PIL import Image as PILImage
import StringIO

# ROS imports
import rospy
import rospkg
from sensor_msgs.msg import CompressedImage, Image
from cv_bridge import CvBridge, CvBridgeError

# Kamera Imports
from custom_msgs.srv import RequestImageMetadata, RequestCompressedImageView, \
    RequestImageView
from sensor_models.image_renderer import warp_perspective


rospack = rospkg.RosPack()


# Instantiate CvBridge
bridge = CvBridge()


class ImageReceiver(object):
    """Receives ROS images and makes them availabe.

    Provides windowed or reduced-resolution image access over network.

    When a request for imagery is made, it is not returned until a new image is
    received by this node.

    """
    def __init__(self, image_topic, image_callback, decoding='passthrough'):
        """
        :param image_topic: Topic providing Image messages.
        :type image_topic: str

        :param image_callback: Function to call with the Image message after
            received by this class. This class manages access to the latest-
            received images, but the image messages can be passed upstream via
            this function.
        :type image_callback: function

        :param decoding: Desired encoding output to pass to ROS cv bridget.
        :type decoding: {'passthrough', 'bgr8', 'rgb8', 'mono8', 'mono16'}

        """
        self.image_lock = threading.RLock()
        self.image_topic = image_topic
        self.image_callback = image_callback

        # The service topic providing access to windowed or reduced-resolution
        # imagery.
        self.image_service_topic = '%s/view_service' % image_topic

        # The service topic providing metadata for the raw imagery stored on
        # this server.
        self.metadata_service_topic = '%s/metadata_service' % image_topic

        self.last_image_msg = None

        rospy.loginfo('Subscribing to Image topic \'%s\''
                      % image_topic)

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

        self.image_metadata_pub = rospy.Publisher('%s_nopixels' % image_topic,
                                                  Image, queue_size=10)


        if self.image_service_topic is not None:
            srv_topic = self.image_service_topic
            rospy.loginfo('Creating RequestImageView service to provide '
                          '\'%s\' image views on topic \'%s\'' %
                          (image_topic, srv_topic))
            rospy.Service(srv_topic, RequestImageView,
                          lambda req: self.image_patch_service_request(req,
                                                                       False,
                                                                       decoding,
                                                                       image_topic))

        if self.image_service_topic is not None:
            srv_topic = '%s/compressed' % self.image_service_topic
            rospy.loginfo('Creating RequestCompressedImageView service to '
                          'provide \'%s\' image views on topic \'%s\'' %
                          (image_topic, srv_topic))
            rospy.Service(srv_topic,
                          RequestCompressedImageView,
                          lambda req: self.image_patch_service_request(req,
                                                                       True,
                                                                       decoding,
                                                                       image_topic))

        if self.metadata_service_topic is not None:
            rospy.loginfo('Creating RequestImageMetadata service to '
                          'provide \'%s\' image metadata via '
                          'RequestImageMetadata on topic \'%s\'' %
                          (image_topic, self.metadata_service_topic))
        rospy.Service(self.metadata_service_topic, RequestImageMetadata,
                      self.metadata_service_topic_request)

    def image_callback_(self, msg, topic):
        with self.image_lock:
            rospy.loginfo('Received Image message on topic \'%s\'' % topic)
            self.last_image_msg = msg

        # Publish an no-pixels version of this image on the image metadata
        # topic.
        msg_metadata = Image()
        msg_metadata.header = msg.header

        # Check whether message is Image/CompressedImage.
        if hasattr(msg, 'format'):
            # Message is CompressedImage.
            pass
        else:
            # Message is Image.
            msg_metadata.height = msg.height
            msg_metadata.width = msg.width
            msg_metadata.encoding = msg.encoding
            msg_metadata.is_bigendian = msg.is_bigendian
            msg_metadata.step = msg.step

        self.image_metadata_pub.publish(msg_metadata)

    def image_patch_service_request(self, req, compress,
                                    decoding='passthrough', topic=None):
        """
        see custom_msgs/srv/RequestImagePatches.srv for more details.

        :param decoding: Desired encoding output to pass to ROS cv bridget.
        :type decoding: {'passthrough', 'bgr8', 'rgb8', 'mono8', 'mono16'}

        """
        rospy.loginfo('Requesting%s \'%s\' image view of size %i x %i that is '
                      'newer than time %0.4f' %
                      (' compressed' if compress else '', topic,
                       req.output_width, req.output_height, req.newer_than_time))

        with self.image_lock:
            if self.last_image_msg is None:
                if compress:
                    return (False, CompressedImage())
                else:
                    return (False, Image())

        if True:
            rospy.loginfo('Image time: %0.3f' %
                          self.last_image_msg.header.stamp.to_sec())
            rospy.loginfo('newer_than_time: %0.3f' %
                          req.newer_than_time)
            rospy.loginfo('Image newness: %0.3f' %
                          (self.last_image_msg.header.stamp.to_sec() - req.newer_than_time))

        # Get the image.
        while True:
            with self.image_lock:
                if self.last_image_msg.header.stamp.to_sec() > req.newer_than_time:
                    img_msg = self.last_image_msg
                    break

            # Wait a little so this loop doesn't dominate CPU.
            time.sleep(0.01)

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

        if req.interpolation == 4:
            flags = cv2.INTER_LANCZOS4 | cv2.WARP_INVERSE_MAP
        elif req.interpolation == 3:
            flags = cv2.INTER_CUBIC | cv2.WARP_INVERSE_MAP
        elif req.interpolation == 2:
            flags = cv2.INTER_AREA | cv2.WARP_INVERSE_MAP
        elif req.interpolation == 1:
            flags = cv2.INTER_LINEAR | cv2.WARP_INVERSE_MAP
        else:
            flags = cv2.INTER_NEAREST | cv2.WARP_INVERSE_MAP

        dsize = (req.output_width, req.output_height)

        homography = np.reshape(req.homography, (3,3)).astype(np.float32)
        image = cv2.warpPerspective(image, homography, dsize=dsize,
                                    flags=flags)

        if compress:
            compressed_msg = CompressedImage()
            compressed_msg.format = "jpeg"
            compressed_msg.data = np.array(cv2.imencode('.jpg',
                                                        image)[1]).tostring()
            compressed_msg.header = img_msg.header
            return (True, compressed_msg)
        else:
            new_image_msg = bridge.cv2_to_imgmsg(image, encoding="rgb8")
            new_image_msg.header = img_msg.header
            return (True, new_image_msg)

    def metadata_service_topic_request(self, req):
        """
        see custom_msgs/srv/RequestImagePatches.srv for more details.

        """
        # We want to return the next image received.
        rospy.loginfo('Metadata requested regarding imagery on topic \'%s\''
                      % self.image_topic)
        with self.image_lock:
            img_msg = self.last_image_msg

            if img_msg is None:
                #print('No image has been received yet, responding with None')
                return (False, 0, 0, '')

            # Check whether received message is Image/CompressedImage.
            if hasattr(img_msg, 'format'):
                # Message is CompressedImage.
                sio = StringIO.StringIO(img_msg.data)
                im = PILImage.open(sio)
                image = np.array(im, copy=False)
                print((True, image.shape[0], image.shape[1], 'rgb8'))
                return (True, image.shape[0], image.shape[1], 'rgb8')
            else:
                return (True, img_msg.height, img_msg.width, img_msg.encoding)


class ImageContainer(object):
    def __init__(self, img, t):
        self.img = img
        self.t = t


class ImageProcessQueue(object):
    def __init__(self, image_topic):
        self.image_topic = image_topic
        self.image_lock = threading.RLock()
        self.images = []

        # Start thread.
        thread = threading.Thread(target=self.daemon, args=())
        thread.daemon = True
        thread.start()

    def add_ros_image(self, msg):
        #print('Got ROS image')
        pass

    def add_image(self, img):
        pass

    def daemon(self):
        ros_rate = rospy.Rate(100)
        while True:
            with self.image_lock:
                pass

            ros_rate.sleep()


class Nexus(object):
    def __init__(self):
        self.image_queues = {}

    def add_image_topic(self, image_topic, decoding='passthrough'):
        """
        :param image_topic: Topic providing Image messages.
        :type image_topic: str

        :param decoding: Desired encoding output to pass to ROS cv bridget.
        :type decoding: {'passthrough', 'bgr8', 'rgb8', 'mono8', 'mono16'}

        """
        if image_topic in self.image_queues:
            raise Exception('Image topic \'%s\' was already added' %
                            image_topic)

        image_queue = ImageProcessQueue(image_topic)
        self.image_queues[image_topic] = image_queue
        ImageReceiver(image_topic, image_queue.add_ros_image, decoding)


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

    def get_param(name, required=False, default=None, loginfo=True):
        """Get ROS parameter with 'name'.

        """
        ros_param = '%s/%s' % (node_name, name)
        param = rospy.get_param(ros_param, default)

        if isinstance(param, str) and param.lower() == 'none':
            param = None

        if required and param is None:
            raise Exception('Parameter \'%s\' must be set' % ros_param)

        if loginfo:
            rospy.loginfo('%s : %s' % (name, str(param)))

        return param

    # -------------------------- Read Parameters -----------------------------
    rgb_image_topic = get_param('rgb_image_topic', True, None, True)
    rgb_analytics_image_topic = get_param('rgb_analytics_image_topic',
                                          False, None, True)
    # ------------------------------------------------------------------------

    nexus = Nexus()

    nexus.add_image_topic(rgb_image_topic, 'rgb8')

    if rgb_analytics_image_topic is not None:
        nexus.add_image_topic(rgb_analytics_image_topic, 'rgb8')

    rospy.spin()

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