#!/usr/bin/env python

# Copyright (C) 2014 Fetch Robotics Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Author: Michael Ferguson

import os
import yaml
import rospy
import rosbag
from sensor_msgs.msg import JointState

if __name__=='__main__':
    # Enumerate files in directory, sort by name
    files = [f for f in os.listdir(".") if f.endswith(".yaml")]
    files.sort()

    # Bag for output
    bag = rosbag.Bag('calibration_poses.bag', 'w')

    # TODO: load this from config file
    names = {}
    names["head_traj_controller"] = ["head_pan_joint", "head_tilt_joint"]
    names["arm_controller"] = ["arm_lift_joint", "arm_shoulder_pan_joint", "arm_shoulder_lift_joint", "arm_elbow_flex_joint", "arm_wrist_flex_joint", "arm_wrist_roll_joint"]

    # Turn each yaml file into joint states in bag
    for filename in files:
        # Parse yaml
        imported = yaml.load(open(filename).read())
        # Convert
        try:
            js = JointState()
            for command in imported["joint_commands"]:
                js.name += names[command["controller"]]
                js.position += command["segments"][-1]["positions"]
            bag.write('calibration_joint_states', js)
            print("Imported %s" % filename)
        except KeyError:
            # must not be a valid config
            pass

    bag.close()