if 'google.colab' in str(get_ipython()):
  from google.colab import drive
  drive.mount('/content/drive')

!pip3 install rosbag --extra-index-url https://rospypi.github.io/simple/

import rosbag
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from scipy import fftpack

def read_slam_data(path):

  bag = rosbag.Bag(path)

  conel_x, conel_y = [], []
  coner_x, coner_y = [], []
  car_x, car_y = [], []
  time = []

  for topic, msg, t in bag.read_messages(topics=['/cone_right', '/cone_left', '/pose_stamped']):
      time.append(t)
      if topic == "/cone_right":
          for pose in msg.poses:
              coner_x.append(pose.position.x)
              coner_y.append(pose.position.y)
      elif topic == "/cone_left":
          for pose in msg.poses:
              conel_x.append(pose.position.x)
              conel_y.append(pose.position.y)
      elif topic == "/pose_stamped":
          car_x.append(msg.pose.position.x)
          car_y.append(msg.pose.position.y)

  length = min(len(time), len(car_x), len(car_y))
  pose_dict = {'time': time[:length], 'x': car_x[:length], 'y': car_y[:length]}
  cone_l_dict = {'x': conel_x, 'y': conel_y}
  cone_r_dict = {'x': coner_x, 'y': coner_y}
  car_pose = pd.DataFrame(pose_dict)
  cone_l = pd.DataFrame(cone_l_dict)
  cone_r = pd.DataFrame(cone_r_dict)
  bag.close()

  return car_pose, cone_l.sort_values(by=['x']), cone_r.sort_values(by=['x'])

if 'google.colab' in str(get_ipython()):
  slam_bag_file_path = '/content/drive/MyDrive/SA/slam.bag'
else:
  slam_bag_file_path = '../data/slam.bag'

car_pose, cone_l, cone_r = read_slam_data(slam_bag_file_path)