import rosbag
import sensor_msgs.point_cloud2 as pc2
import cv2
import argparse
from sensor_msgs.msg import Image
from sensor_msgs.msg import PointCloud2
from cv_bridge import CvBridge, CvBridgeError
import rospy
import os
import keyboard
import pypcd
# import Queue
import Queue
import numpy as np  
IMAGE,TIME = None, None
# global recordLidar
recordLidar=False
q=Queue.Queue()
parser = argparse.ArgumentParser(description="listen to a topic")
parser.add_argument('--cameratopicName', type=str, required=True, help='name of camera topic')
parser.add_argument('--lidartopicName', type=str, required=True, help='name of lidar topic')
parser.add_argument('--savePath', type=str, required=True, help="image save path")
args = parser.parse_args()
TOPICNAME1, TOPICNAME2,SAVEPATH = args.cameratopicName, args.lidartopicName,args.savePath

num=0
Path=None
def image_callback(image):
    bridge = CvBridge()
    cv_image = bridge.imgmsg_to_cv2(image, "bgr8")
    IMAGE = cv_image
    TIME = "%.6f" % image.header.stamp.to_sec()
    # return cv_image
    cv2.imshow("img", cv_image)
    k = cv2.waitKey(1) & 0xFF
    global recordLidar,num,Path
    if k == ord('s'):
        if not os.path.exists(SAVEPATH+"/"+str(num)):
            os.mkdir(SAVEPATH+"/"+str(num))
            Path=SAVEPATH+"/"+str(num)
            num=num+1
        recordLidar=True
        cv2.imwrite(Path +"/" + str(num) + ".jpg", IMAGE)
    elif k == ord('q'):
        exit(0)

def lidar_callback(lidar):
    global q
    q.put(lidar)    
    if q.qsize()>5:
        q.get()
    global recordLidar,Path
    if recordLidar==True: 
        for i in range(q.qsize()):
            sublidar=q.get()
            pc = pypcd.PointCloud.from_msg(sublidar)
            filename = "%.6f" % sublidar.header.stamp.to_sec() + '.pcd'
            laser_data_path = os.path.join(Path, filename)    
            pc.save(laser_data_path)  
            # data = pc2.read_points(sublidar)
            # points = np.array(list(data), dtype=np.float32) 
            # filename = "%.6f" % sublidar.header.stamp.to_sec() + '.pcd'
            # file_path = os.path.join(SAVEPATH, filename)           
            # points.tofile(file_path)
            recordLidar=False    

   

rospy.init_node("imagelidarSaver")
rospy.Subscriber(TOPICNAME1, Image, image_callback)
rospy.Subscriber(TOPICNAME2, PointCloud2, lidar_callback)

rospy.spin()