from rosbag import Bag
from cv_bridge import CvBridge
from cv2 import imwrite
import json
import argparse
import os



def main():

    # --------------------------------------------------------
    # Arguments
    # --------------------------------------------------------
    
    ap = argparse.ArgumentParser()
    ap.add_argument("-dataset_name", "--dataset_name",
                    help="Dataset name. Datasets must be inside hpe/images/.",
                    type=str, default='sim_moving')
    ap.add_argument("-bagfile", "--bagfile",
                    help="Path to bagfile",
                    type=str, default='/home/daniela/bagfiles/hpe_ok.bag')

    args = vars(ap.parse_args())

    # --------------------------------------------------------
    # Loading files, and verifications configurations
    # --------------------------------------------------------
    cv_bridge = CvBridge()

    bag = Bag('/home/daniela/bagfiles/hpe_moving.bag')

    dataset_folder = '../../images/' + args['dataset_name'] + '/'
    if not os.path.exists(dataset_folder):
        raise Exception("The dataset folder does not exist!")


    ground_truth_path = "ground_truth.txt"

    with open(ground_truth_path, "r") as fp:
        gt_dic = json.load(fp)

    sync_dic = {}

    topic_rgb = '/camera_1/rgb/image_raw'
    window_name = 'image'
    i = 0
    n = 0
    for topic_rgb, msg, t in bag.read_messages(topics=[topic_rgb]):
        img = cv_bridge.imgmsg_to_cv2(msg)
        p = dataset_folder + 'camera_1/' + str(i) + '.png'

        value = int(str(t))
        timestamps = gt_dic.keys()
        timestamps = [eval(i) for i in timestamps]
        res = min(timestamps, key=lambda x: abs(value - x))
        sync_dic[i] = gt_dic[str(res)]


        imwrite(p, img)
        i = i + 1

    with open(dataset_folder + "ground_truth.txt", "w") as fp:
        json.dump(sync_dic, fp)

    topic_rgb = '/camera_2/rgb/image_raw'
    window_name = 'image'
    i = 0
    n = 0
    for topic_rgb, msg, t in bag.read_messages(topics=[topic_rgb]):
        img = cv_bridge.imgmsg_to_cv2(msg)
        p = dataset_folder + 'camera_2/' + str(i) + '.png'
        imwrite(p, img)
        i = i + 1

    topic_rgb = '/camera_3/rgb/image_raw'
    window_name = 'image'
    i = 0
    n = 0
    for topic_rgb, msg, t in bag.read_messages(topics=[topic_rgb]):
        img = cv_bridge.imgmsg_to_cv2(msg)
        p = dataset_folder + 'camera_3/' + str(i) + '.png'
        imwrite(p, img)
        i = i + 1


    topic_rgb = '/camera_4/rgb/image_raw'
    window_name = 'image'
    i = 0
    n = 0
    for topic_rgb, msg, t in bag.read_messages(topics=[topic_rgb]):
        img = cv_bridge.imgmsg_to_cv2(msg)
        p = dataset_folder + 'camera_4/' + str(i) + '.png'
        imwrite(p, img)
        i = i + 1




if __name__ == "__main__":
    main()