# coding: utf-8
# !/usr/bin/env python2
import os
import rosbag
import csv
import math
import rospy
import sys
import time
import numpy as np
from datetime import datetime
import argparse
import pandas as pd

# from nav_msgs.msg import OccupancyGrid


global_height, global_costmap, global_origin, global_resolution = None, None, None, None


def calculate_start_index(g_h, local_origin, global_origin, resolution):
    print(global_origin.position.x, global_origin.position.y)
    x = g_h - 100 - int((local_origin.position.y - global_origin.position.y) / resolution) - 2
    y = int((local_origin.position.x - global_origin.position.x) / resolution) + 1
    return x, y


def compare_costmaps(local_costmap, global_costmap, start_x, start_y):
    dynamic_obstacles = []
    # print("global_costmap.shape[0]:", global_costmap.shape[0])
    # print("global_costmap.shape[1]:", global_costmap.shape[1])
    for i in range(100):
        for j in range(100):
            global_x = start_x + i
            global_y = start_y + j
            if global_x < 0 or global_y < 0 or global_x >= global_costmap.shape[0] or global_y >= global_costmap.shape[
                1]:
                continue
            if global_costmap[global_x][global_y] < 50 and local_costmap[i][j] >= 50:
                dynamic_obstacles.append((global_x, global_y))
    return dynamic_obstacles


def process_global_rosbag(bagfile, global_topic):
    global global_costmap, global_origin, global_resolution, global_height
    with rosbag.Bag(bagfile, 'r') as bag:
        for g_topic, g_msg, g_t in bag.read_messages(topics=global_topic):
            global_costmap = np.array(g_msg.data).reshape(g_msg.info.height, g_msg.info.width)
            global_costmap = np.flipud(global_costmap)  # Flip the global costmap vertically
            global_origin = g_msg.info.origin
            global_resolution = g_msg.info.resolution
            global_height = g_msg.info.height
            break


def process_local_rosbag(bagfile, local_topic):
    global global_costmap, global_origin, global_resolution, global_height
    all_dynamic_obstacles = []
    detected_obstacles = set()
    is_first_frame = True

    local_messages = []
    with rosbag.Bag(bagfile, 'r') as bag:
        start_time = bag.get_start_time()
        end_time = bag.get_end_time()
        total_time = end_time - start_time
        print("total_time:", total_time)
        for topic, msg, t in bag.read_messages(topics=local_topic):
            local_messages.append((msg, t))

        local_interval = total_time / len(local_messages)

        for index, (msg, t) in enumerate(local_messages):
            local_costmap = np.array(msg.data).reshape(msg.info.height, msg.info.width)
            local_costmap = np.flipud(local_costmap)  # Flip local costmap vertically
            local_origin = msg.info.origin
            start_x, start_y = calculate_start_index(global_costmap.shape[0], local_origin, global_origin,
                                                     msg.info.resolution)
            frame_time = index * local_interval
            dynamic_obstacles = compare_costmaps(local_costmap, global_costmap, start_x, start_y)
            if is_first_frame:
                all_dynamic_obstacles.extend([(frame_time, x, y) for x, y in dynamic_obstacles])
                detected_obstacles.update((x, y) for x, y in dynamic_obstacles)
                is_first_frame = False
            else:
                filtered_obstacles = [(frame_time, x, y) for x, y in dynamic_obstacles if
                                      (x, y) not in detected_obstacles]
                all_dynamic_obstacles.extend(filtered_obstacles)
                detected_obstacles.update((x, y) for _, x, y in filtered_obstacles)

    df = pd.DataFrame(all_dynamic_obstacles, columns=['frame_time', 'x', 'y'])

    return df


def merge_obstacles(df, global_height):
    df.sort_values(by=['frame_time', 'x', 'y'], inplace=True)

    merged_time = []
    merged_rows = []
    merged_columns = []
    merged_counts = []
    merge_count = 1

    for i in range(1, len(df)):
        if df.iloc[i]['frame_time'] == df.iloc[i - 1]['frame_time'] and df.iloc[i]['x'] == df.iloc[i - 1]['x'] + 1 and \
                df.iloc[i]['y'] == df.iloc[i - 1]['y']:
            merge_count += 1
        else:
            mid_frame_time = df.iloc[i - 1]['frame_time']
            mid_row = df.iloc[i - (merge_count // 2 + 1)]['x']
            mid_column = df.iloc[i - 1]['y']
            merged_time.append(mid_frame_time)
            merged_rows.append(mid_row)
            merged_columns.append(mid_column)
            merged_counts.append(merge_count)
            merge_count = 1

    mid_frame_time = df.iloc[-1]['frame_time']
    mid_row = df.iloc[-1 - (merge_count // 2)]['x']
    mid_column = df.iloc[-1]['y']
    merged_time.append(mid_frame_time)
    merged_rows.append(mid_row)
    merged_columns.append(mid_column)
    merged_counts.append(merge_count)

    result_df = pd.DataFrame({'Time': merged_time, 'x': merged_rows, 'y': merged_columns, 'Count': merged_counts})
    result_df['pos_x'] = result_df['y'].apply(lambda x: (x - 1) * 0.05 + 0.05 if x % 2 != 0 else (x - 1) * 0.05 + 0.025)
    result_df['pos_y'] = (global_height - result_df['x']) * 0.05 + 0.025  # ????? height
    print("global_height:", global_height)
    result_df['Scaling_x'] = result_df['Count'] * 0.05
    result_df['Scaling_y'] = 0.05
    return result_df


def parsezihao(input_dir, output_dir):
    dic_EgoState = ['Time', 'simTime', 'simFrame', 'posX', 'posY', 'posZ', 'speedX', 'speedY', 'speedZ', 'accelX',
                    'accelY', 'accelZ',
                    'dimX', 'dimY', 'dimZ', 'obstacle', 'traveledDist']
    # dic_DriverCtrl=['Time','simTime','simFrame','tarspeedX','tarspeedY','tarspeedZ','tardimX','tardimY','tardimZ']
    EgoState_file = open(output_dir + "/" + "EgoState_pji.csv", 'w')
    # DriverCtrl_file = open(output_dir + "/"+"DriverCtrl_pji.csv", 'w')
    writer_EgoState = csv.writer(EgoState_file)
    writer_EgoState.writerow(dic_EgoState)
    # writer_DriverCtrl = csv.writer(DriverCtrl_file)
    # writer_DriverCtrl.writerow(dic_DriverCtrl)

    frame_max = sys.maxsize
    count = 1
    with rosbag.Bag(input_dir, 'r') as bag:
        odom_flag = False
        first_message_time = None
        Frame_imu = 1
        Frame_cmd_vel = 1
        obstacle_state = 0
        cur_mileage = ''

        for topic, msg, t in bag.read_messages():  # t代表时间

            if first_message_time is None:
                first_message_time = t
                first_message_time = rospy.Time.to_sec(first_message_time)
                first_message_time = datetime.fromtimestamp(first_message_time)

            if topic == "/obstacle_detection":
                obstacle_state = msg.data
                # print(msg.data)

            if topic == "/sys_info":
                cur_mileage = msg.cur_mileage

            if topic == "/odom":
                odom_flag = True
                posX = msg.pose.pose.position.x
                posY = msg.pose.pose.position.y
                posZ = msg.pose.pose.position.z
                speedX = msg.twist.twist.linear.x * 3.6
                speedY = msg.twist.twist.linear.y * 3.6
                speedZ = msg.twist.twist.linear.z * 3.6
                dimX = msg.twist.twist.angular.x
                dimY = msg.twist.twist.angular.y
                dimZ = msg.twist.twist.angular.z

            if topic == "/imu":

                if odom_flag:
                    accelX = msg.linear_acceleration.x
                    accelY = msg.linear_acceleration.y
                    accelZ = msg.linear_acceleration.z
                    timestamp = rospy.Time.to_sec(t)
                    date_time_imu = datetime.fromtimestamp(timestamp)
                    simtime_imu = (date_time_imu - first_message_time).total_seconds()

                    message_EgoState = [date_time_imu, simtime_imu, Frame_imu, posX, posY, posZ, speedX, speedY, speedZ,
                                        accelX, accelY,
                                        accelZ, dimX, dimY, dimZ, obstacle_state, cur_mileage]

                    writer_EgoState.writerow(message_EgoState)
                    Frame_imu += 1
                else:
                    print('6666')
            '''        
            if topic =='/cmd_vel':
                timestamp = rospy.Time.to_sec(t)
                date_time_cmd_vel = datetime.fromtimestamp(timestamp)
                simtime_cmd_vel=(date_time_cmd_vel-first_message_time).total_seconds()
                message_DriverCtrl =[date_time_cmd_vel,simtime_cmd_vel,Frame_cmd_vel,msg.linear.x*3.6,msg.linear.y*3.6,msg.linear.z*3.6,
                                     msg.angular.x,msg.angular.y,msg.angular.z]
                
                                     
                writer_DriverCtrl.writerow(message_DriverCtrl)
                Frame_cmd_vel+=1
            '''

        EgoState_file.close()
        # DriverCtrl_file.close()


def parsehancheng(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__":

    global_bagfile = sys.argv[1]  # 全局地图的rosbag包
    # global_bagfile = '/media/dell/56FAD34BFAD32651/data/pujin_datareturn/20240428/costmap.bag'
    input_dir = sys.argv[2]  # 触发采集的rosbag包
    # input_dir='/media/dell/56FAD34BFAD32651/data/pujin_datareturn/20240428/2024-04-25-15-58-27_obstacledetection_30.bag'
    bagname = input_dir.split('/')[-1].split('.')[0]
    output_dir = sys.argv[3]  # 输出文件路径
    # output_dir='/media/dell/56FAD34BFAD32651/data/pujin_datareturn/20240428'
    global_topic = '/move_base/global_costmap/costmap'
    local_topic = '/move_base/local_costmap/costmap'

    output_dir = os.path.join(output_dir, bagname)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    '''
    df, global_height = process_rosbag(input_dir, global_topic, local_topic)
    result_df = merge_obstacles(df, global_height)
    result_df.to_csv(os.path.join(output_dir,'merged_obstacles.csv'), index=False)
    parsehancheng(input_dir, output_dir)
    parsezihao(input_dir, output_dir)
    print('successfully analysis '+input_dir)
    '''

    try:
        process_global_rosbag(global_bagfile, global_topic)
        df = process_local_rosbag(input_dir, local_topic)
        result_df = merge_obstacles(df, global_height)
        result_df.to_csv(os.path.join(output_dir, 'merged_obstacles.csv'), index=False)
        parsehancheng(input_dir, output_dir)
        parsezihao(input_dir, output_dir)
        print('successfully analysis ' + input_dir)
    except Exception as e:
        print(e)


def parse(costmap_bag, data_bag, output_dir):
    global_bagfile = costmap_bag  # 全局地图的rosbag包
    # global_bagfile = '/media/dell/56FAD34BFAD32651/data/pujin_datareturn/20240428/costmap.bag'
    input_dir = data_bag  # 触发采集的rosbag包
    # input_dir='/media/dell/56FAD34BFAD32651/data/pujin_datareturn/20240428/2024-04-25-15-58-27_obstacledetection_30.bag'
    bagname = input_dir.split('/')[-1].split('.')[0]
    # output_dir='/media/dell/56FAD34BFAD32651/data/pujin_datareturn/20240428'
    global_topic = '/move_base/global_costmap/costmap'
    local_topic = '/move_base/local_costmap/costmap'

    output_dir = os.path.join(output_dir, bagname)
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    '''
    df, global_height = process_rosbag(input_dir, global_topic, local_topic)
    result_df = merge_obstacles(df, global_height)
    result_df.to_csv(os.path.join(output_dir,'merged_obstacles.csv'), index=False)
    parsehancheng(input_dir, output_dir)
    parsezihao(input_dir, output_dir)
    print('successfully analysis '+input_dir)
    '''

    try:
        process_global_rosbag(global_bagfile, global_topic)
        df = process_local_rosbag(input_dir, local_topic)
        result_df = merge_obstacles(df, global_height)
        result_df.to_csv(os.path.join(output_dir, 'merged_obstacles.csv'), index=False)
        parsehancheng(input_dir, output_dir)
        parsezihao(input_dir, output_dir)
        print('successfully analysis ' + input_dir)
    except Exception as e:
        print(e)
