#!/usr/bin/env python

# Script to record (a subtree of) tf data to the db
# 
# Requires
# - db to be running
# - configuration file containing transforms to save
# - the required transforms to be published by other ros nodes
#
# Config file consists of lines of the form PARENT CHILD, e.g.
# ---
# map base_footprint
# base_footprint camera_rgb_optical_frame
# ---
# 
# would allow keeping track of where the base is and where the head is 
# pointing, in the map frame

import roslib
roslib.load_manifest('mongo_ros')

import rospy
import argparse
import sys
import tf
from tf.msg import tfMessage
import mongo_ros as mr
import geometry_msgs.msg as gm

def parse_transform_file(filename):
    "Given the filename, return list of pairs (PARENT_FRAME, CHILD_FRAME)"
    transforms = []
    with open(filename) as f:
        for line in f:
            tokens = line.split()
            if len(tokens)==0:
                continue
            elif len(tokens)==2:
                transforms.append((tokens[0], tokens[1]))
            else:
                print("Skipping malformed line {0}".format(line))
    return transforms
    

if __name__=='__main__':
    # Parse command line arguments
    argv = rospy.myargv(sys.argv)
    parser = argparse.ArgumentParser(description="Record tf data to the db")
    parser.add_argument('transform_file', help="File containing the transforms"
                        " to save.  It must consist of lines of the form "
                        "'PARENT_FRAME CHILD_FRAME'.")
    parser.add_argument('db_name',
                        help="Name of db to save to within the Mongo instance")
    parser.add_argument('--rate', '-r', default=1.0, type=float,
                        help="Rate (hz) to save the given transforms at")
    parser.add_argument('--host', '-o', default='localhost', help="Host for db server")
    parser.add_argument('--port', '-p', default=27017, help="Port for db server", type=int)

    args = parser.parse_args(argv[1:])
    transforms = parse_transform_file(args.transform_file)
    sys.stderr.write("Saving these transforms at {0} hz\n".format(args.rate))
    for tr in transforms:
        sys.stderr.write("  {0} <- {1}\n".format(tr[0], tr[1]))

    # Set up ros and db connections
    rospy.init_node('save_transforms', anonymous=True)
    listener = tf.TransformListener()
    coll = mr.MessageCollection(args.db_name, 'tf', tfMessage,
                                indexes=['stamp'], db_host=args.host,
                                db_port=args.port)
    sys.stderr.write("Node initialized\n")
    
    # Main loop
    rate = rospy.Rate(args.rate)
    while not rospy.is_shutdown():
        # Iterate over transforms we care about
        for parent, child in transforms:
            try:
                # Lookup the transform
                t = listener.getLatestCommonTime(child, parent)
                (trans, (x,y,z,w)) = listener.lookupTransform(parent, child, t)
                # TODO we should probably check that the transform is recent

                # Make the transform message
                tr = gm.TransformStamped()
                tr.header.stamp = t
                tr.header.frame_id = parent
                tr.child_frame_id = child
                tr.transform.translation = gm.Vector3(*trans)
                tr.transform.rotation = gm.Quaternion(x=x, y=y, z=z, w=w)
            
                # Save it
                coll.insert(tfMessage([tr]), {'stamp': t.to_sec()})
                
            except tf.Exception as e:
                rospy.logwarn("Skipping transform due to exception {0}".\
                              format(e))

        rate.sleep()
            

            
            
