#!/usr/bin/env python
# -*- coding: utf-8 -*-
import sys
import os
import subprocess
import multiprocessing
import rosbag


def convert_bag_to_pcd(input_dir, output_dir):
    command = 'rosrun pcl_ros bag_to_pcd {} /scan_map_icp_amcl_node/scan_point_transformed {}'.format(input_dir,
                                                                                                      output_dir)
    process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, stderr = process.communicate()


def calculate_frame_rate(rosbag_path):
    with rosbag.Bag(rosbag_path, 'r') as bag:
        num_count = 0
        for topic, msg, t in bag.read_messages():
            if topic == "/scan_map_icp_amcl_node/scan_point_transformed":
                num_count += 1
        duration = bag.get_end_time() - bag.get_start_time()
        frame_rate = num_count / float(duration) if duration > 0 else 0
    return frame_rate


def save_frame_rate(lidar_output_base, frame_rate):
    frame_rate_file = os.path.join(lidar_output_base, 'frame_rate.txt')
    with open(frame_rate_file, 'w') as file:
        file.write(str(frame_rate))


def process_file(file):
    root, filename = file
    inputFilePath = os.path.join(root, filename)
    ascii_folder = os.path.join(root, '../pcd_ascii')
    try:
        os.makedirs(ascii_folder)
    except Exception:
        print '创建目录失败 %s' % ascii_folder
        pass
    tempFilePath = os.path.join(ascii_folder, filename)

    command = "pcl_convert_pcd_ascii_binary {} {} 0".format(inputFilePath, tempFilePath)
    os.system(command)


def main():
    rosbag_path = sys.argv[1]
    output_base_path = sys.argv[2]
    bag_base_name = os.path.basename(rosbag_path).split('.')[0]
    lidar_output_base = os.path.join(output_base_path, bag_base_name + '_pcd_lidar')
    pcd_folder_path = os.path.join(lidar_output_base, 'pcd')

    if not os.path.exists(pcd_folder_path):
        os.makedirs(pcd_folder_path)

    # Convert ROS bag to PCD
    convert_bag_to_pcd(rosbag_path, pcd_folder_path)

    # Calculate and save frame rate
    frame_rate = calculate_frame_rate(rosbag_path)
    save_frame_rate(lidar_output_base, frame_rate)

    # Process files in parallel
    files_to_process = []
    for root, dirs, files in os.walk(pcd_folder_path):
        for file in files:
            if file.endswith(".pcd"):
                files_to_process.append((root, file))

    num_processes = 4
    pool = multiprocessing.Pool(num_processes)
    pool.map(process_file, files_to_process)
    pool.close()
    pool.join()


if __name__ == '__main__':
    main()
