#! /usr/bin/env python

from __future__ import print_function
import sys
import picamera
import picamera.array
import cv2
import numpy as np
import rospy
from rospy.numpy_msg import numpy_msg
from sensor_msgs.msg import CompressedImage
from tabletop.msg import MotionVectorArray

class FrameAnalyzer(picamera.array.PiAnalysisOutput):

    def __init__(self, camera, pub, frame_id, size):
        self.pub = pub
        self.frame_id = frame_id

    def write(self, frame):
        now = rospy.Time.now()
        rospy.logdebug('Analyzing frame, size=%s', len(frame))

        msg = CompressedImage()
        msg.header.stamp = now
        msg.header.frame_id = self.frame_id
        msg.format = 'jpeg'
        msg.data = bytes(frame)
        self.pub.publish(msg)
        rospy.logdebug('Finished analyzing frame')

class MotionAnalyzer(picamera.array.PiMotionAnalysis):

    def __init__(self, camera, pub, size):
        super(MotionAnalyzer, self).__init__(camera)
        self.pub = pub
        self.size = size

    def analyze(self, vectors):
        now = rospy.Time.now()
        rospy.logdebug('Analyzing vectors, size=%s', vectors.shape)

        h, w = vectors.shape
        vectors = vectors.reshape(h*w)

        msg = MotionVectorArray()
        msg.header.stamp = now
        msg.image_width = self.size[0]
        msg.image_height = self.size[1]
        msg.width = w
        msg.height = h
        msg.x = vectors['x']
        msg.y = vectors['y']
        msg.sad = vectors['sad']
        self.pub.publish(msg) 
        rospy.logdebug('Finished analyzing vectors')

class VideoPublisher:

    def __init__(self):
        pass

    def run(self):
        rospy.init_node('publish_video')

        frame_rate = rospy.get_param('~frame_rate', 10)
        image_width = rospy.get_param('~image_width', 640)
        image_height = rospy.get_param('~image_height', 480)
        vector_width = rospy.get_param('~vector_width', 640)
        vector_height = rospy.get_param('~vector_height', 480)
        frame_id = rospy.get_param('~frame_id', 'camera_link')

        image_pub = rospy.Publisher('image/compressed', CompressedImage,
                                    queue_size=10)
        vector_pub = rospy.Publisher('motion_vectors',
                                     numpy_msg(MotionVectorArray),
                                     queue_size=10)

        with picamera.PiCamera() as camera:
            image_size = (image_width, image_height)
            vector_size = (vector_width, vector_height)
            camera.resolution = vector_size
            camera.framerate = frame_rate

            with FrameAnalyzer(camera, image_pub, frame_id, image_size) \
                 as frame_analyzer:
                with MotionAnalyzer(camera, vector_pub, vector_size) \
                     as motion_analyzer:
                    camera.start_recording(frame_analyzer, 'mjpeg',
                        resize=image_size)
                    camera.start_recording('/dev/null', 'h264',
                        splitter_port=2, intra_period=0,
                        motion_output=motion_analyzer)
                    try:
                        while not rospy.is_shutdown():
                            camera.wait_recording(1)
                    except:
                        print('Unexpected error: {0}'.format(sys.exc_info()[0]))
                        raise
                    finally:
                        camera.stop_recording(splitter_port=2)
                        camera.stop_recording()

if __name__ == "__main__":
    try:
        VideoPublisher().run()

    except rospy.ROSInterruptException:
        pass
