#! /usr/bin/env python
# -*- coding: utf-8 -*-
'''
Created on 2022/8/10 12:06:24

@author: LiuKuan
@copyright: Apache License, Version 2.0
'''
import rosbag
# ===============================================================================
'''遍历/读取 rosbag数据
'''
# ===============================================================================

__all__ = ['rosbag_iter', 'rosbag_iter_img', 'rosbag_iter_pcd', 'save_pcd', 'save_img']

def rosbag_iter(rosbag_name, topic_name = None):
    with rosbag.Bag(rosbag_name, "r") as bag:  # 要读取的bag文件；
        for topic, msg, t in bag.read_messages():
            if topic_name is None:
                yield topic, msg, t
            elif topic_name == topic:
                yield topic, msg, t

def rosbag_iter_img(rosbag_name, topic_name = None):
    from cv_bridge import CvBridge
    bridge = CvBridge()
    for topic, msg, t in rosbag_iter(rosbag_name, topic_name):
        cv_image = bridge.imgmsg_to_cv2(msg, "bgr8")
        yield cv_image, msg.header.stamp.to_sec()

def rosbag_iter_pcd(rosbag_name, topic_name = None):
    import sensor_msgs.point_cloud2 as pc2
    import numpy as np

    for topic, msg, t in rosbag_iter(rosbag_name, topic_name):
        points_data = np.array(list(pc2.read_points(msg)))
        lidar = list(np.delete(points_data, np.where(np.isnan(points_data))[0], axis=0))
        yield lidar, msg.header.stamp.to_sec()

def save_pcd(lidar, filename, binary=True):
    if binary:
        _save_pcd_binary(lidar, filename)
    else:
        _save_pcd_ascii(lidar, filename)

def _save_pcd_binary(lidar, filename):
    PCD_BINARY_TEMPLATE = """VERSION 0.7
FIELDS x y z intensity
SIZE 4 4 4 4
TYPE F F F F
COUNT 1 1 1 1
WIDTH {}
HEIGHT 1
VIEWPOINT 0 0 0 1 0 0 0
POINTS {}
DATA binary
"""

    import struct
    sz = len(lidar)
    with open(filename, 'wb') as f:
        header = PCD_BINARY_TEMPLATE.format(sz, sz)
        f.write(header.encode())
        for pi in lidar:
            h = struct.pack('ffff', pi[0], pi[1], pi[2], pi[3])
            f.write(h)

def _save_pcd_ascii(lidar, filename):
    PCD_ASCII_TEMPLATE = """VERSION 0.7
FIELDS x y z intensity
SIZE 4 4 4 2
TYPE F F F U
COUNT 1 1 1 1
WIDTH {}
HEIGHT 1
VIEWPOINT 0 0 0 1 0 0 0
POINTS {}
DATA ascii
"""
    sz = len(lidar)
    with open(filename, 'w') as f:
        header = PCD_ASCII_TEMPLATE.format(sz, sz)
        f.write(header)
        for pi in lidar:
            f.write(' '.join([str(i) for i in pi]) + '\n')

def save_img(img, filename):
    import os, cv2
    name, ext = os.path.splitext(filename)
    retval, buf = cv2.imencode(ext, img)
    return buf.tofile(filename)
