# import bagpy
from bagpy import bagreader
import pandas as pd
import numpy as np
import pickle as pkl
import os
import rosbag

class RobotTraj():
    def __init__(self) -> None:
        self.desired_topic = '/anna/end_effector/states'
        self.other_topics =  ['/anna/joint/states', '/anna/keyframe/states']

        self.users =['user2','user4','user7','user8','user10','user14','user16','user18','user20'] # narration
        self.users += ['user3','user5','user6','user9','user11','user12','user15','user17','user19'] # natural

        self.tasks = ['cutting', 'box']
        self.demo_types = ['video', 'kt']

        self.data_dir = '/Volumes/Seagate Portable Drive/audio_study/kinesthetic/'

        self.path_len = {}
        self.path_len_time = {}

    def compute_path_len(self):
        # iterate through users
        for user in self.users:
            for task in self.tasks:
                for demo_type in self.demo_types:
                    key = user+'_'+task+'_'+demo_type

                    if key not in self.path_len_time:
                        self.path_len_time[key] = {}

                    print(user,task,demo_type)
                    bagfiles = []
                    bag_path = os.path.join(self.data_dir,user,task,demo_type) 
                    for dirname, dirs, files in os.walk(bag_path):
                        for filename in files:
                            fn,ex = os.path.splitext(filename)
                            if ex == '.bag':
                                bagfiles.append(filename)
                    bagfiles.sort(reverse=True)
                    bag = os.path.join(bag_path,bagfiles[0])
                    bag_audio = rosbag.Bag(bag) 

                    total_dist = 0.
                    i = 0
                    for idx, data in enumerate(bag_audio.read_messages(topics=[self.desired_topic])):
                        topic, msg, t = data

                        # print(msg)
                        if i==0:
                            x,y,z = msg.pose.position.x,msg.pose.position.y,msg.pose.position.z
                            old = np.array((x,y,z))
                            t_start = t.secs*pow(10,9) + t.nsecs
                            curr_time = 0
                        else:
                            x,y,z = msg.pose.position.x,msg.pose.position.y,msg.pose.position.z
                            newl = np.array((x,y,z))
                            total_dist += np.linalg.norm(newl-old)
                            t_new = t.secs*pow(10,9) + t.nsecs
                            curr_time = t_new-t_start

                            old = np.copy(newl)

                        # print(curr_time*pow(10,-9))
                        curr_time = curr_time*pow(10,-9)
                        # print(round(curr_time,1))
                        
                        # save distance upto every 0.1 second
                        self.path_len_time[key][round(curr_time,1)] = total_dist
                        i+=1

                    print(f'Total distance traveled by the arm is {total_dist}')
                    
                    self.path_len[key] = total_dist

    def save_path_len(self):
        self.compute_path_len()
        filename = '../../data/robot_path_len.pkl'
        with open(filename, 'wb') as handle:
            pkl.dump(self.path_len, handle, protocol=pkl.HIGHEST_PROTOCOL)

        filename = '../../data/robot_path_len_time.pkl'
        with open(filename, 'wb') as handle:
            pkl.dump(self.path_len_time, handle, protocol=pkl.HIGHEST_PROTOCOL)



def main():
    traj = RobotTraj()
    traj.save_path_len()

if __name__ == "__main__":
    main()
# import bagpy
from bagpy import bagreader
import pandas as pd
import numpy as np
import pickle as pkl
import os
import rosbag

class RobotTraj():
    def __init__(self) -> None:
        self.desired_topic = '/anna/end_effector/states'
        self.other_topics =  ['/anna/joint/states', '/anna/keyframe/states']

        self.users =['user2','user4','user7','user8','user10','user14','user16','user18','user20'] # narration
        self.users += ['user3','user5','user6','user9','user11','user12','user15','user17','user19'] # natural

        self.tasks = ['cutting', 'box']
        self.demo_types = ['video', 'kt']

        self.data_dir = '/Volumes/Seagate Portable Drive/audio_study/kinesthetic/'

        self.path_len = {}
        self.path_len_time = {}

    def compute_path_len(self):
        # iterate through users
        for user in self.users:
            for task in self.tasks:
                for demo_type in self.demo_types:
                    key = user+'_'+task+'_'+demo_type

                    if key not in self.path_len_time:
                        self.path_len_time[key] = {}

                    print(user,task,demo_type)
                    bagfiles = []
                    bag_path = os.path.join(self.data_dir,user,task,demo_type) 
                    for dirname, dirs, files in os.walk(bag_path):
                        for filename in files:
                            fn,ex = os.path.splitext(filename)
                            if ex == '.bag':
                                bagfiles.append(filename)
                    bagfiles.sort(reverse=True)
                    bag = os.path.join(bag_path,bagfiles[0])
                    bag_audio = rosbag.Bag(bag) 

                    total_dist = 0.
                    i = 0
                    for idx, data in enumerate(bag_audio.read_messages(topics=[self.desired_topic])):
                        topic, msg, t = data

                        # print(msg)
                        if i==0:
                            x,y,z = msg.pose.position.x,msg.pose.position.y,msg.pose.position.z
                            old = np.array((x,y,z))
                            t_start = t.secs*pow(10,9) + t.nsecs
                            curr_time = 0
                        else:
                            x,y,z = msg.pose.position.x,msg.pose.position.y,msg.pose.position.z
                            newl = np.array((x,y,z))
                            total_dist += np.linalg.norm(newl-old)
                            t_new = t.secs*pow(10,9) + t.nsecs
                            curr_time = t_new-t_start

                            old = np.copy(newl)

                        # print(curr_time*pow(10,-9))
                        curr_time = curr_time*pow(10,-9)
                        # print(round(curr_time,1))
                        
                        # save distance upto every 0.1 second
                        self.path_len_time[key][round(curr_time,1)] = total_dist
                        i+=1

                    print(f'Total distance traveled by the arm is {total_dist}')
                    
                    self.path_len[key] = total_dist

    def save_path_len(self):
        self.compute_path_len()
        filename = '../../data/robot_path_len.pkl'
        with open(filename, 'wb') as handle:
            pkl.dump(self.path_len, handle, protocol=pkl.HIGHEST_PROTOCOL)

        filename = '../../data/robot_path_len_time.pkl'
        with open(filename, 'wb') as handle:
            pkl.dump(self.path_len_time, handle, protocol=pkl.HIGHEST_PROTOCOL)



def main():
    traj = RobotTraj()
    traj.save_path_len()

if __name__ == "__main__":
    main()
# import bagpy # narration # natural # iterate through users # print(msg) # print(curr_time*pow(10,-9)) # print(round(curr_time,1)) # save distance upto every 0.1 second