#!/usr/bin/env python

# Copyright (C) 2022 Michael Ferguson
#
# 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.

import rospy
import rosbag
import sys
import yaml

if __name__=='__main__':
    if len(sys.argv) < 2:
        print("Usage:")
        print("  convert_ros1_bag_to_yaml <bagfile> <default_features>")
        print("")
        print("  This script is run on ROS1 to export a ROS1 bagfile of")
        print("  capture poses into a YAML file that can be used in ROS2")
        exit(-1)

    # Poses to export as YAML
    poses = list()

    # Extract list of features
    features = sys.argv[2:]

    # Convert bagfile
    bag_name = sys.argv[1]
    print("Migrating {}".format(bag_name))
    bag = rosbag.Bag(bag_name)
    for topic, msg, t in bag.read_messages(topics=["calibration_joint_states"]):
        pose = dict()
        try:
            # Load robot_calibration_msgs/CaptureConfig
            pose["joints"] = list(msg.joint_states.name)
            pose["positions"] = list(msg.joint_states.position)
            pose["features"] = list(msg.features)
        except AttributeError:
            # Load older sensor_msgs/JointState type bagfile
            pose["joints"] = list(msg.name)
            pose["positions"] = list(msg.position)
            if features:
                pose["features"] = list(features)
        poses.append(pose)

    # Save yaml file
    with open("capture_poses.yaml", "w") as file:
        yaml.dump(poses, file)
