#!/usr/bin/env python

# Script to save an entire ros bag into a db where you can get random access to it

import roslib
roslib.load_manifest('mongo_ros')
import argparse
import mongo_ros as mr
import rosbag as rb
import rospy
import sys
from mongo_ros import util
import math

def remove_leading_slash(n):
    return n[1:] if n.startswith('/') else n

def is_power_of_two(x):
    y = math.log(x, 2)
    d = abs(y-round(y))
    return d<1e-3

class Recorder(object):
    """
    Class that holds the state of the overall recording process.

    The main mutable state is the map from topic name to the corresponding
    MessageCollection, class object, and other info.
    """

    def __init__(self, topics, db, bag):
        self.topics = topics or None
        self.db = db
        self.bag = bag
        self.collections = {}

    def record_all(self):
        "Save all messages to the db"
        b = rb.Bag(self.bag)

        # Iterate over the messages
        ind = 0
        for top, tup, t in b.read_messages(topics=self.topics, raw=True):
            topic = remove_leading_slash(top)
            
            # If we haven't seen this topic before,
            # do any necessary imports, and set up the MessageCollection object
            if topic not in self.collections:
                pkg, msg = tup[0].split('/')
                cls = util.load_message_class(pkg, msg)
                coll = mr.MessageCollection(self.db, topic, cls)
                self.collections[topic] = (cls, coll)
            cls, coll= self.collections[topic]

            m = cls()
            m.deserialize(tup[1])
            coll.insert(m, {})

            ind += 1
            if ind % 10 == 0:
                sys.stderr.write("\r{0} messages processed".format(ind))
            

        


############################################################
# Main entry point
############################################################

def main(argv=sys.argv):
    rospy.init_node('bag_to_db')
    argv = rospy.myargv(argv)
    parser = argparse.ArgumentParser(description="Save bag messages to mongodb")
    parser.add_argument('db_name', help="Name of db")
    parser.add_argument('bag_file', help="Bag file")
    parser.add_argument('--all_topics', '-a', help="Record all topics",
                        action='store_true', default=False)
    parser.add_argument('--topic', '-t', help="To record topics T1, T2, T3, do"
                        " -t T1 -t T2 -t T3", default=[], action='append')
    args = parser.parse_args(argv[1:])

        
    print "Saving messages from {0.db_name} to {0.bag_file}".format(args)
    if args.topic:
        print "Recording topics ", args.topic
    elif args.all_topics:
        print "Recording all topics"
    else:
        print "Must use specify topics or use --all_topics"
        return 1

    r = Recorder(args.topic, args.db_name, args.bag_file)
    r.record_all()

if __name__=='__main__':
    sys.exit(main())
