#!/usr/bin/python
import rospy
import sys
import struct
import argparse
import time
import numpy
import traceback
from sensor_msgs.msg import Image
from snark.imaging import cv_image
import cv2
import rosbag

help_description='listen to a ros topic for raw image and write it in cv format to stdout'
help_example="""
example
  read from published ros topic and output image
    ros-to-image /camera/rgb/image_color | cv-cat "timestamp;view;null"

  read from rosbags and output image
    ros-to-image /camera/rgb/image_color /log/*.bag | cv-cat "timestamp;view;null"

"""
parser=argparse.ArgumentParser(description=help_description,epilog=help_example,formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument('topic',help='topic to listen to')
parser.add_argument('bags', metavar='BAGFILE', nargs='*', help='read messages from bag files')
parser.add_argument('--flush',help='flush stdout after each image frame write',action='store_true')
args=parser.parse_args()

encoding_map={
    'rgb8': cv2.CV_8UC3, 
    'rgba8': cv2.CV_8UC4,
    "rgb16": cv2.CV_16UC3,
    "rgba16": cv2.CV_16UC4,
    "bgr8": cv2.CV_8UC3,
    "bgra8": cv2.CV_8UC4,
    "bgr16": cv2.CV_16UC3,
    "bgra16": cv2.CV_16UC4,
    "mono8": cv2.CV_8UC1,
    "mono16": cv2.CV_16UC1,
    "8UC1": cv2.CV_8UC1,
    "8UC2": cv2.CV_8UC2,
    "8UC3": cv2.CV_8UC3,
    "8UC4": cv2.CV_8UC4,
    "8SC1": cv2.CV_8SC1,
    "8SC2": cv2.CV_8SC2,
    "8SC3": cv2.CV_8SC3,
    "8SC4": cv2.CV_8SC4,
    "16UC1": cv2.CV_16UC1,
    "16UC2": cv2.CV_16UC2,
    "16UC3": cv2.CV_16UC3,
    "16UC4": cv2.CV_16UC4,
    "16SC1": cv2.CV_16SC1,
    "16SC2": cv2.CV_16SC2,
    "16SC3": cv2.CV_16SC3,
    "16SC4": cv2.CV_16SC4,
    "32SC1": cv2.CV_32SC1,
    "32SC2": cv2.CV_32SC2,
    "32SC3": cv2.CV_32SC3,
    "32SC4": cv2.CV_32SC4,
    "32FC1": cv2.CV_32FC1,
    "32FC2": cv2.CV_32FC2,
    "32FC3": cv2.CV_32FC3,
    "32FC4": cv2.CV_32FC4,
    "64FC1": cv2.CV_64FC1,
    "64FC2": cv2.CV_64FC2,
    "64FC3": cv2.CV_64FC3,
    "64FC4": cv2.CV_64FC4,
    "bayer_rggb8": cv2.CV_8UC4,
    "bayer_bggr8": cv2.CV_8UC4,
    "bayer_gbrg8": cv2.CV_8UC4,
    "bayer_grbg8": cv2.CV_8UC4,
    "bayer_rggb16": cv2.CV_16UC4,
    "bayer_bggr16": cv2.CV_16UC4,
    "bayer_gbrg16": cv2.CV_16UC4,
    "bayer_grbg16": cv2.CV_16UC4
}

shutdown=False

def message_to_image(msg):
    global shutdown
    global args
    try:
        header=numpy.zeros((),dtype=cv_image.image.header_dtype)
        t=msg.header.stamp
        header['time']=t.secs*1000000+t.nsecs/1000
        header['rows']=msg.height
        header['cols']=msg.width
        header['type']=encoding_map[msg.encoding]
        header.tofile(sys.stdout)
        #t=msg.header.stamp
        #ts=struct.pack('q',t.secs*1000000+t.nsecs/1000)
        #sys.stdout.write(ts)
        #sys.stdout.write(struct.pack('I',msg.height))
        #sys.stdout.write(struct.pack('I',msg.width))
        #sys.stdout.write(struct.pack('I',encoding_map[msg.encoding]))
        sys.stdout.write(msg.data)
        if args.flush:
            sys.stdout.flush()
        
    except SystemExit:
        shutdown=True
    except:
        shutdown=True
        sys.stderr.write("exception!\n")
        traceback.print_exc()
        
if len(args.bags):
    for bagfile in args.bags:
        with rosbag.Bag(bagfile) as bag:
            for topic, msg, t in bag.read_messages(topics=args.topic):
                if shutdown: exit(0)
                message_to_image(msg)
else:
    rospy.init_node('listener',anonymous=True,disable_signals=True)
    subscription = rospy.Subscriber(args.topic, Image, message_to_image)
    while not shutdown:
        time.sleep(0.001)
    rospy.signal_shutdown('shutdown')
