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


allNum = []
path = sys.path[0]+"/"
image_raw = ""
image_raw_png = ""
image_raw_comp = ""
image_raw_comp_png = ""
pointCloud_path = ""


def createFolder():
        #if not os.path.exists(image_raw+num+'/'):
         #       os.mkdir(image_raw+num+'/')
        #if not os.path.exists(image_raw_png+num+'/'):
         #       os.mkdir(image_raw_png+num+'/')
        #if not os.path.exists(image_raw_comp+num+'/'):
         #       os.mkdir(image_raw_comp+num+'/')
        #if not os.path.exists(image_raw_comp_png+num+'/'):
         #       os.mkdir(image_raw_comp_png+num+'/')
        if not os.path.exists(pointCloud_path+'/'):
                os.mkdir(pointCloud_path+'/')


def saveImage(num, topic, msg, bridge):
        date_1 = date_2 = []
        if not num in allNum:
                allNum.append(num)
                createFolder(num)
        if "compressed" in topic:
                try:
                        cv_img_comp = bridge.compressed_imgmsg_to_cv2(
                            msg, 'bgr8')
                except CvBridgeError as e:
                        print e
                timestr = "%.6f" % msg.header.stamp.to_sec()
               # image_name = timestr+'.bmp'
                #image_name_png = timestr+'.png'
                #cv2.imwrite(image_raw_comp+num+'/'+image_name, cv_img_comp)
                #cv2.imwrite(image_raw_comp_png+num+'/' +image_name_png, cv_img_comp)
        else:
                try:
                        cv_img = bridge.imgmsg_to_cv2(msg, 'bgr8')
                except CvBridgeError as e:
                        print e
                timestr = "%.6f" % msg.header.stamp.to_sec()
                #image_name = timestr+'.bmp'
                #image_name_png = timestr+'.png'
                #cv2.imwrite(image_raw+num+'/'+image_name, cv_img)
                #cv2.imwrite(image_raw_png+num+'/'+image_name_png, cv_img)


def main(argv):
        if len(argv) < 2:
                print "please give a bag path"
                sys.exit(1)
	#print path
	bag_name = argv[1].split('/')[-1].split('.')[0]
        path = argv[1].split('.bag')[0]
       # path = path + bag_name
        if not os.path.exists(path):
                os.mkdirs(path+"/pandar")

        global image_raw
        image_raw = path+'/im'
        global image_raw_png
        image_raw_png = path+'/imPng'
        global image_raw_comp
        image_raw_comp = path+'/imComp'
        global image_raw_comp_png
        image_raw_comp_png = path+'/imCompPng'
        global pointCloud_path
        #pointCloud_path = path+'/pandar'
        #bridge = CvBridge()
	#if not os.path.exists(pointCloud_path+'/'):
	#	os.mkdir(pointCloud_path+'/')
        print "****************begin***************"
        with rosbag.Bag(argv[1], 'r') as bag:
		with open(path + "/" + bag_name + ".txt","w") as w :
		        for topic, msg, t in bag.read_messages():
		                if "/image_raw" in topic and (len(topic) == 20 or len(topic) == 31):
		                        if "1" in topic:
		                                saveImage("1", topic, msg, bridge)
		                        elif "2" in topic:
		                                saveImage("2", topic, msg, bridge)
		                        elif "3" in topic:
		                                saveImage("3", topic, msg, bridge)
		                elif topic == "/pandar":
					#timestr = "%.6f" % msg.header.stamp.to_sec()
		                        tstr = list(str(t))
		                        tstr.insert(10,'.')
		                        tname = ''.join(tstr)
					w.write(tname+"\n")
		                        laser_data_name = tname + ".pcd"
		                        laser_data_path = os.path.join(pointCloud_path,laser_data_name)
		                        pc = pypcd.PointCloud.from_msg(msg)
		                        pc.save(laser_data_path)
                        
        bag.close()
        print "****************done****************"

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