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

import sys
import time


def parsebag(input_dir, output_dir):
    dic_object_detection = ['Time', 'FrameID', 'HeadingAngle', 'ID', 'East', 'North', 'position_z', 'altitude',
                            'AbsVel', 'Type', 'width', 'length', 'height']
    object_detection_file = open(output_dir + "/" + "pos.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(topics=['/cicv_location', '/tpperception']):  # t代表时间

            if topic == "/cicv_location":  # 100hz
                hasLoc = True
                ego_speed = math.sqrt(msg.velocity_x ** 2 + msg.velocity_y ** 2 + msg.velocity_z ** 2)
                message_location = [str(t)[:-6], framenum, msg.yaw, '-1', msg.position_x, msg.position_y,
                                    msg.position_z, msg.altitude, ego_speed, '2', '0', '0', '0']

                global pos_x, pos_y, pos_z
                pos_x = msg.position_x
                pos_y = msg.position_y
                pos_z = msg.position_z
                egocar_yaw = msg.yaw

            if topic == "/tpperception":  # 10hz
                output_type = ' '
                if not hasLoc:
                    print("666")
                    continue
                # global  pos_x ,pos_y,pos_z
                m = msg.header.sequence_num
                if m < frame_max:
                    frame_max = m

                msg_l = len(msg.objs)
                # print(msg_l)
                # seq=msg.header.sequence_num-frame_max+1
                for i in range(msg_l):

                    obj_x = msg.objs[i].xabs
                    obj_y = msg.objs[i].yabs
                    obj_z = msg.objs[i].z + pos_z
                    # print(msg.objs[i].type)

                    if msg.objs[i].type == 2:
                        output_type = '2'  # 乘用车/轿车
                    elif msg.objs[i].type == 3:
                        output_type = '4'  # 货车/皮卡
                    elif msg.objs[i].type == 1:
                        output_type = '7'  # 行人
                    elif msg.objs[i].type == 4:
                        output_type = '8'  # 两轮车
                    elif msg.objs[i].type == 0:
                        output_type = '10'  # other
                    elif msg.objs[i].type == 5:
                        output_type = '103'  # other
                    message_obj = [message_location[0], framenum, msg.objs[i].heading, str(msg.objs[i].id), obj_x,
                                   obj_y, obj_z, message_location[-6], msg.objs[i].speed, output_type,
                                   msg.objs[i].width, msg.objs[i].length, msg.objs[i].height]
                    writer_object_detection.writerow(message_obj)
                writer_object_detection.writerow(message_location)
                framenum += 1
                # hasLoc = False
                # data_read_flag=False

        # driving_status_file.close()
        object_detection_file.close()


def parse(input_dir, output_dir):
    bagname = input_dir.split('/')[-1].split('.')[0]
    output_dir = os.path.join(output_dir, bagname)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    try:
        parsebag(input_dir, output_dir)
    except Exception as e:
        print(e)
