import rosbag
import os
import sys
import cv2
import numpy as np
#import pypcd
import sensor_msgs.point_cloud2 as pc2
from cv_bridge import CvBridge
from sensor_msgs.msg import Image
from cv_bridge import CvBridgeError

def check_bag_time(bag_path,date):
        time0 = time1 = time2 = 0
        with rosbag.Bag(bag_path, 'r') as bag:
                i = 0
                for topic, msg, t in bag.read_messages():
                        if i == 0 :
                                time0 = msg.header.stamp.to_sec()
                        if i == 1 :
                                time1 = msg.header.stamp.to_sec()
                        if i == 2 :
                                time2 = msg.header.stamp.to_sec()
                                break
                if (time0 == time1) and (time2 - time1) == 0.1:
                        date.append(time0 - 0.1) 
                else:
                        date.append(time0) 
                      

def save2png(path,date):
        #  check_bag_time(bag_path,date)
         bridge = CvBridge()         
        #  path = bag_path.split('.')[0]
         if not os.path.exists(path+'/camera_0/'):
                os.makedirs(path+'/camera_0/')
         if not os.path.exists(path+'/camera_2/'):
                os.makedirs(path+'/camera_2/')
         if not os.path.exists(path+'/camera_1/'):
                os.makedirs(path+'/camera_1/')
         with rosbag.Bag(path+".bag", 'r') as bag:
                for topic, msg, t in bag.read_messages():
                        if "/image_raw" in topic :
                                if "left" in topic :
                                        try:
                                                cv_img = bridge.imgmsg_to_cv2(msg, 'bgr8')
                                        except CvBridgeError as e:
                                                print e
                                        timestr = "%.6f" % msg.header.stamp.to_sec()
                                        cv2.imwrite(path+'/camera_0/'+timestr+".png", cv_img)
                                elif "right" in topic :
                                        try:
                                                cv_img = bridge.imgmsg_to_cv2(msg, 'bgr8')
                                        except CvBridgeError as e:
                                                print e
                                        timestr = "%.6f" % msg.header.stamp.to_sec()
                                        cv2.imwrite(path+'/camera_1/'+timestr+".png", cv_img)
                                elif "middle" in topic :
                                        try:
                                                cv_img = bridge.imgmsg_to_cv2(msg, 'bgr8')
                                        except CvBridgeError as e:
                                                print e
                                        timestr = "%.6f" % msg.header.stamp.to_sec()
                                        cv2.imwrite(path+'/camera_2/'+timestr+".png", cv_img)
                                        
         bag.close()

def getCommonTime(path,date_1,date_2):
    common = []
    if date_2.count(date_1[0]):             # date_2 ahead of date_1
        num = date_1.count(date_2[-1])      # match the time from the last one of date_2 to the current place of date_1
        for i in range(num,len(date_2)):
            common.append(date_2[i])
            i = i + 1
    else:
        num = date_2.count(date_1[-1])
        for i in range(0,num):
            common.append(num,len(date_1))
            i = i + 1
    with open(path+'/times.txt','w') as w:
        for i in range(0,len(common)):
            w.write("%.6f" % (common[i])+'\n')
    w.close()

def main(argv):
        if len(argv) < 1:
                print "please give a bag path and a time txt path"
                sys.exit(1)
        # date_1 = date_2 = []
        date = []
        path = argv[1].split('.bag')[0]
        # path = save2png(argv[1],date_1)
        save2png(path,date)
        # getCommonTime(path,date_1,date_2)



# function begin
if __name__ == "__main__":
        main(sys.argv)