# coding: utf-8
#!/usr/bin/env python2
import os
import rosbag
import csv
import math 

import sys
import time
import numpy as np

def parsebag(input_dir, output_dir):   
    def quaternion_to_euler(x, y, z, w):
        # 将四元数归一化
        try:
            length = np.sqrt(x**2 + y**2 + z**2 + w**2)
            x /= length
            y /= length
            z /= length
            w /= length
        
            # 计算欧拉角
            #roll = np.arctan2(2*(w*x + y*z), 1 - 2*(x**2 + y**2))
            #pitch = np.arcsin(2*(w*y - z*x))
            yaw = np.arctan2(2*(w*z + x*y), 1 - 2*(y**2 + z**2))  
            return  yaw
        except :
            return 0


    dic_object_detection = ['Time','FrameID','HeadingAngle','X', 'Y' ,'Z']
    object_detection_file = open(output_dir + "/"+"pos_pji.csv", 'w')
    writer_object_detection = csv.writer(object_detection_file)
    writer_object_detection.writerow(dic_object_detection)


    frame_max=sys.maxsize
    with rosbag.Bag(input_dir ,'r') as bag:
        #flag=False
        framenum = 1
        #hasLoc = False
        for topic,msg,t in bag.read_messages():    #t代表时间
            
            if topic == "/amcl_pose":#100hz  /odom
                poseX=msg.pose.pose.position.x
                poseY=msg.pose.pose.position.y
                poseZ=msg.pose.pose.position.z
                orientationX=msg.pose.pose.orientation.x
                orientationY=msg.pose.pose.orientation.y
                orientationZ=msg.pose.pose.orientation.z
                orientationW=msg.pose.pose.orientation.w
                egoyaw=quaternion_to_euler(orientationX,orientationY,orientationZ,orientationW)
                
                

                message_location =[str(t)[:-6],framenum,egoyaw,poseX,poseY,poseZ]

       
                
                
                  
                writer_object_detection.writerow(message_location)
                framenum+=1
                #hasLoc = False
                #data_read_flag=False
                
    # driving_status_file.close()
        object_detection_file.close()

if __name__ == "__main__":
#    print("Input the path of file...."+'\n')

   #input_dir = sys.argv[1]
   input_dir='/media/dell/BFAC-F22B/data/pujin_datareturn/pjirobot_20231207_2023-12-07-07-18-46_54.bag'
   bagname=input_dir.split('/')[-1].split('.')[0]
   #output_dir = sys.argv[2]
   output_dir='/media/dell/BFAC-F22B/data/pujin_datareturn'

   output_dir=os.path.join(output_dir, bagname)
   if not os.path.exists(output_dir):
       os.makedirs(output_dir)
   #parsebag(input_dir, output_dir)

   try:
       parsebag(input_dir, output_dir)
       print('successfully analysis '+input_dir)
   except Exception as e:
       print(e)