from rosbag import Bag
import numpy as np
import time
import re
import json
from gazebo_msgs.msg import LinkStates

bag = Bag('/home/daniela/bagfiles/hpe_moving.bag')

topic_link_states = '/gazebo/link_states'
dic = {}
for topic_link_states, msg, t in bag.read_messages(topics=[topic_link_states]):
    # print(msg)
    t = int(str(t))
    dic[t] = {}
    i = 0
    for name in msg.name:
        if "actor::" in name:
            joint = name[len("actor::"):]
            # print(joint)
            if joint != "actor_pose":
                dic[t][joint] = {}
                dic[t][joint]['pose'] = {}
                dic[t][joint]['pose']['x'] = msg.pose[i].position.x
                dic[t][joint]['pose']['y'] = msg.pose[i].position.y
                dic[t][joint]['pose']['z'] = msg.pose[i].position.z
                dic[t][joint]['orientation'] = {}
                dic[t][joint]['orientation']['x'] = msg.pose[i].orientation.x
                dic[t][joint]['orientation']['y'] = msg.pose[i].orientation.y
                dic[t][joint]['orientation']['z'] = msg.pose[i].orientation.z
                dic[t][joint]['orientation']['w'] = msg.pose[i].orientation.w

        i = i + 1
    # print(dic)
    # exit(0)

with open("ground_truth.txt", "w") as fp:
    # Load the dictionary from the file
    json.dump(dic, fp)
#   res_key, res_val = min(d.items(), key=lambda x: abs(value - x[1])) find closest number in dic


# joints

# Hips
# LHipJoint
# LeftUpLeg
# LeftLeg
# LeftFoot
# LeftToeBase
# RHipJoint
# RightUpLeg
# RightLeg
# RightFoot
# RightToeBase
# LowerBack
# Spine
# Spine1
# Neck
# Neck1
# Head
# LeftShoulder
# LeftArm
# LeftForeArm
# LeftHand
# LeftFingerBase
# LeftHandIndex1
# LThumb
# RightShoulder
# RightArm
# RightForeArm
# RightHand
# RightFingerBase
# RightHandIndex1
# RThumb