#!/usr/bin/env python
import sys
import roslib;
import rospy
import rosbag
from rospy import rostime
import CommonFun
import argparse
import os
import yaml


def parse_args():
    parser = argparse.ArgumentParser(
        prog = 'bagmerge.py',
        description='Merges two bagfiles.')
    parser.add_argument('-o', type=str, help='name of the output file', 
        default = None, metavar = "output_file")
    parser.add_argument('-t', type=str, help='topics which should be merged to the main bag', 
        default = None, metavar = "topics")
    parser.add_argument('-i', help='reindex bagfile', 
        default = False, action="store_true")
    parser.add_argument('main_bagfile', type=str, help='path to a bagfile, which will be the main bagfile')
    parser.add_argument('bagfile', type=str, help='path to a bagfile which should be merged to the main bagfile')
    args = parser.parse_args()
    return args

def get_next(bag_iter, reindex = False, 
        main_start_time = None, start_time = None, 
        topics = None):
    try:
        result = bag_iter.next()
        if topics != None:
            while not result[0] in topics:
                result = bag_iter.next()
        if reindex:
            return (result[0], result[1], 
                result[2] - start_time + main_start_time)
        return result
    except StopIteration:
        return None

def merge_bag(main_bagfile, bagfile, outfile = None,sensor_topic= None, imu_topic = None, 
        reindex = True):
    #get min and max time in bagfile
    main_limits = get_limits(main_bagfile,sensor_topic)
    limits = get_limits(bagfile,imu_topic)
    #check output file
    if outfile == None:
        pattern = main_bagfile + "_merged_%i.bag"
        outfile = main_bagfile + "_merged.bag"
        index = 0
        while (os.path.exists(outfile)):
            outfile = pattern%index
            index += 1
    #output some information
    print ("merge bag %s in %s"%(bagfile, main_bagfile))
    print ("topics filter: ", sensor_topic,imu_topic)
    print ("writing to %s."%outfile)
    #merge bagfile
    outbag = rosbag.Bag(outfile, 'w')
    main_bag = rosbag.Bag(main_bagfile).__iter__()
    bag = rosbag.Bag(bagfile).__iter__()
    main_next = get_next(main_bag)
    next = get_next(bag, reindex, main_limits[0], limits[0], imu_topic)
    try:
        while main_next != None or next !=None:
        
            if main_next == None:
                
                if next[2].to_sec()<=main_limits[1].to_sec()+0.1 and next[2].to_sec()>=main_limits[0].to_sec():
                    outbag.write(next[0], next[1], next[2])
                elif next[2].to_sec()>main_limits[1].to_sec():
                    break
                next = get_next(bag, reindex, main_limits[0], limits[0], imu_topic)
            #elif next == None:
                #outbag.write(main_next[0], main_next[1], main_next[2])
                #main_next = get_next(main_bag)
            #elif next[2].to_sec()<=main_limits[1].to_sec() and next[2].to_sec()>=main_limits[0].to_sec():
                #outbag.write(next[0], next[1], next[2])
                #next = get_next(bag, reindex, main_limits[0], limits[0], imu_topic)                

            else:
                
                outbag.write(main_next[0], main_next[1], main_next[2])
                main_next = get_next(main_bag)
    finally:
        outbag.close()

def get_limits(bagfile,sensor_topic):
    print ("Determine start and end index of %s..."%bagfile)
    end_time = None
    start_time = None

    for topic, msg, t in rosbag.Bag(bagfile).read_messages():
        if topic==sensor_topic:
            t=msg.header.stamp
            if start_time == None or t < start_time:
                start_time = t
            if end_time == None or t > end_time:
                end_time = t
    return (start_time, end_time)
    
def main():

    optFile = os.getcwd() + "/mergeBag.yaml"
    with open(optFile) as f:
        opt = yaml.safe_load(f)

    mainBagPath = opt["main_bag_path"]
    imuBagPath = opt["imu_bag_path"]
    outBagPath = opt["out_bag_path"]
    sensor_topic=opt["sensor_topic"]

    mainPath = mainBagPath
    imuPath = imuBagPath
    outPath=outBagPath

    mainFiles = []
    imuFiles=[]
    CommonFun.iterate_file_s(mainPath,mainFiles,".bag")
    CommonFun.iterate_file_s(imuPath,imuFiles,".bag")
    for main_bagfile in mainFiles:
        imubagPath=str(imuFiles[0])
        mainbagPath=mainPath+"/"+str(main_bagfile)
        imubagPath=imuPath+"/"+str(imubagPath)
        outbagPath=outPath+"/"+str(main_bagfile)
        merge_bag(mainbagPath, imubagPath, outbagPath,sensor_topic,"/imu_raw",False)


if __name__ == "__main__":
    main()