#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Created on Wed Nov  1 13:38:28 2023

@author: dell
"""
import subprocess
from subprocess import Popen, PIPE
import os
import sys
import math
import cv2
import numpy as np
import shutil
import rosbag


def draw_points(pic, filename):
    f = open(filename, "r")
    ff = f.readlines()[11:]

    for i in range(0, len(ff)):
        point = ff[i].split()[:4]
        point_x = float(point[0])
        point_y = -float(point[1])

        if math.isnan(point_x):
            pass
        else:
            cv2.circle(pic, (int(200 - point_x * 10), int(100 - point_y * 10)), 1, (250, 250, 250), 2)


def rotx(t):
    """Rotation about the x-axis."""
    c = np.cos(t)
    s = np.sin(t)
    return np.array([[1, 0, 0],
                     [0, c, -s],
                     [0, s, c]])


def roty(t):
    """Rotation about the y-axis."""
    c = np.cos(t)
    s = np.sin(t)
    return np.array([[c, 0, s],
                     [0, 1, 0],
                     [-s, 0, c]])


def rotz(t):
    """Rotation about the z-axis."""
    c = np.cos(t)
    s = np.sin(t)
    return np.array([[c, -s, 0],
                     [s, c, 0],
                     [0, 0, 1]])


def transform_from_rot_trans(R, t):
    """Transforation matrix from rotation matrix and translation vector."""
    R = R.reshape(3, 3)
    t = t.reshape(3, 1)
    return np.vstack((np.hstack([R, t]), [0, 0, 0, 1]))


def view(rx, ry, rz, tx, ty, tz):
    mx = rotx(rx)
    my = roty(ry)
    mz = rotz(rz)
    rotation = np.dot(mz, np.dot(my, mx))
    translation = np.array([tx, ty, tz])
    return transform_from_rot_trans(rotation, translation)


# if __name__ == '__main__':
#     input_dir = sys.argv[1]
#     # input_dir='/media/dell/BFAC-F22B/data/jinlong_datareturn/test_1116/jinlong_data_merge_2023-12-13-05-59-34_TTC_7.bag'
#     output_dir = sys.argv[2]
#     # output_dir = '/media/dell/新加卷/样例及代码/金龙车/样例'
#     a = view(0, 30, 0, -10, 0, 5)
#     colormap = np.array([[128, 130, 120],
#                          [235, 0, 205],
#                          [0, 215, 0],
#                          [235, 155, 0]]) / 255.0

#     bag_name = input_dir.split('/')[-1].split('.')[0]
#     output_dir = os.path.join(output_dir, bag_name + '_pcd_forwardlook' + '/pcd')
#     if not os.path.exists(output_dir):
#         os.makedirs(output_dir)

#     with rosbag.Bag(input_dir, 'r') as bag:
#         num_count = 0
#         for topic, msg, t in bag.read_messages():
#             if topic == "/points_concat":  # /camera/depth/points   # /scan_map_icp_amcl_node/scan_point_transformed
#                 num_count += 1

#     bagtime = int(input_dir.split('/')[-1].split('.')[0].split('_')[-1])
#     hz = str(float(num_count) / bagtime)

#     #######################解析rosbag中的点云point_concat话题，获得原始pcd文件#######################
#     command = ['rosrun', 'pcl_ros', 'bag_to_pcd'] + [input_dir] + ['/points_concat'] + [output_dir]
#     command = ' '.join(command)
#     os.system(command)
#     file1 = os.path.join(output_dir[0:-4], 'pcd_ascii')
#     file2 = os.path.join(output_dir[0:-4], 'jpg')
#     if not os.path.exists(file1):
#         os.makedirs(file1)
#     if not os.path.exists(file2):
#         os.makedirs(file2)

#     #######################将原始pcd文件转化成pcd-ascii格式,并生成jpg文件#######################
#     for root, dirs, files in os.walk(output_dir[0:-4]):

#         # find .pcd files and do the transfer twice
#         for file in files:
#             if "pcd" == root[-3:] and ".pcd" == file[-4:]:

#                 inputFilePath = root + "/" + file
#                 tempFilePath = root[:-3] + "pcd_ascii/" + file
#                 outputFilePath = root[:-3] + "jpg/" + file[:-3] + "jpg"

#                 # transfer pcd_utf-8 to pcd_ascii
#                 command = "pcl_convert_pcd_ascii_binary " + \
#                           inputFilePath + \
#                           " " + \
#                           tempFilePath + \
#                           " 0"
#                 # print(command)
#                 os.system(command)

#                 # make sure the temp file exist, invalid input file may cause this issue
#                 if not os.path.exists(tempFilePath):
#                     continue

#                 # transfer pcd_ascii to jpg
#                 # pic = np.zeros([200, 400, 3])  #pic = np.zeros([800, 1600, 3])
#                 # draw_points(pic, tempFilePath)
#                 # cv2.imwrite(outputFilePath, pic)
#                 # print("jpg generated: ", outputFilePath) 
#                 # pcd_path = "D:/文件/2-课题/科技部3.2/影子模式/pcd/data/pcd_ascii/1666841252.226625000.pcd"
#                 pic = np.zeros([800, 1600, 3])
#                 f = open(tempFilePath, "r")
#                 ff = f.readlines()[11:]

#                 for i in range(0, len(ff)):
#                     point = ff[i].split()[:4]
#                     point_x = float(point[0])
#                     point_y = float(point[1])
#                     point_z = float(point[2])
#                     points_intensity = float(point[2])
#                     t = np.dot(np.linalg.inv(a), np.array([point_x, point_y, point_z, 1]))
#                     x = t[0]
#                     y = t[1]

#                     if math.isnan(point_x):
#                         pass
#                     else:
#                         cv2.circle(pic, (int(800 - y * 40), int(400 - x * 40)), 0,
#                                    colormap[int(points_intensity) % colormap.shape[0]] * 255, 0)

#                 cv2.imwrite(outputFilePath, pic)
#             else:
#                 break

#     #######################将转化的点云jpg合成视频#######################
#     jpg_list = os.listdir(file2)
#     if not jpg_list == []:
#         command = ['ffmpeg', '-f', 'image2', '-r', hz, '-pattern_type', 'glob', '-i'] + ['"jpg/*.jpg"'] + ['-tag:v',
#                                                                                                            'avc1',
#                                                                                                            '-y'] + [
#                       'pcd_forwardlook.mp4']
#         result_string = " ".join(command)
#         p = Popen(result_string, shell=True, cwd=output_dir[0:-4] + '/', stdout=subprocess.PIPE, stderr=subprocess.PIPE)
#         p.wait()
#     shutil.rmtree(file1)
#     shutil.rmtree(file2)
#     shutil.rmtree(output_dir)


def parse(input_bag_file, output_mp4_dir):
    input_dir = input_bag_file
    # input_dir='/media/dell/BFAC-F22B/bizhang_2023-11-15-10-35-21.bag'
    output_dir = output_mp4_dir
    # output_dir='/media/dell/BFAC-F22B'
    # output_dir = '/media/dell/新加卷/样例及代码/金龙车/样例'
    a = view(0, 30, 0, -10, 0, 5)
    colormap = np.array([[128, 130, 120],
                         [235, 0, 205],
                         [0, 215, 0],
                         [235, 155, 0]]) / 255.0

    bag_name = input_dir.split('/')[-1].split('.')[0]
    output_dir = os.path.join(output_dir, bag_name + '_pcd_forwardlook' + '/pcd')
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    with rosbag.Bag(input_dir, 'r') as bag:
        num_count = 0
        for topic, msg, t in bag.read_messages():
            if topic == "/points_concat":  # /camera/depth/points   # /scan_map_icp_amcl_node/scan_point_transformed
                num_count += 1

    bagtime = int(input_dir.split('/')[-1].split('.')[0].split('_')[-1])
    hz = str(float(num_count) / bagtime)

    #######################解析rosbag中的点云point_concat话题，获得原始pcd文件#######################
    command = ['rosrun', 'pcl_ros', 'bag_to_pcd'] + [input_dir] + ['/points_concat'] + [output_dir]
    command = ' '.join(command)
    os.system(command)
    file1 = os.path.join(output_dir[0:-4], 'pcd_ascii')
    file2 = os.path.join(output_dir[0:-4], 'jpg')
    if not os.path.exists(file1):
        os.makedirs(file1)
    if not os.path.exists(file2):
        os.makedirs(file2)

    #######################将原始pcd文件转化成pcd-ascii格式,并生成jpg文件#######################
    for root, dirs, files in os.walk(output_dir[0:-4]):

        # find .pcd files and do the transfer twice
        for file in files:
            if "pcd" == root[-3:] and ".pcd" == file[-4:]:

                inputFilePath = root + "/" + file
                tempFilePath = root[:-3] + "pcd_ascii/" + file
                outputFilePath = root[:-3] + "jpg/" + file[:-3] + "jpg"

                # transfer pcd_utf-8 to pcd_ascii
                command = "pcl_convert_pcd_ascii_binary " + \
                          inputFilePath + \
                          " " + \
                          tempFilePath + \
                          " 0"
                # print(command)
                os.system(command)

                # make sure the temp file exist, invalid input file may cause this issue
                if not os.path.exists(tempFilePath):
                    continue

                # transfer pcd_ascii to jpg
                # pic = np.zeros([200, 400, 3])  #pic = np.zeros([800, 1600, 3])
                # draw_points(pic, tempFilePath)
                # cv2.imwrite(outputFilePath, pic)
                # print("jpg generated: ", outputFilePath)
                # pcd_path = "D:/文件/2-课题/科技部3.2/影子模式/pcd/data/pcd_ascii/1666841252.226625000.pcd"
                pic = np.zeros([800, 1600, 3])
                f = open(tempFilePath, "r")
                ff = f.readlines()[11:]

                for i in range(0, len(ff)):
                    point = ff[i].split()[:4]
                    point_x = float(point[0])
                    point_y = float(point[1])
                    point_z = float(point[2])
                    points_intensity = float(point[2])
                    t = np.dot(np.linalg.inv(a), np.array([point_x, point_y, point_z, 1]))
                    x = t[0]
                    y = t[1]

                    if math.isnan(point_x):
                        pass
                    else:
                        cv2.circle(pic, (int(800 - y * 40), int(400 - x * 40)), 0,
                                   colormap[int(points_intensity) % colormap.shape[0]] * 255, 0)

                cv2.imwrite(outputFilePath, pic)
            else:
                break

    #######################将转化的点云jpg合成视频#######################
    jpg_list = os.listdir(file2)
    if not jpg_list == []:
        command = ['ffmpeg', '-f', 'image2', '-r', hz, '-pattern_type', 'glob', '-i'] + ['"jpg/*.jpg"'] + ['-tag:v',
                                                                                                           'avc1',
                                                                                                           '-y'] + [
                      'pcd_forwardlook.mp4']
        result_string = " ".join(command)
        p = Popen(result_string, shell=True, cwd=output_dir[0:-4] + '/', stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        p.wait()
    shutil.rmtree(file1)
    shutil.rmtree(file2)
    shutil.rmtree(output_dir)
    return output_dir[0:-4]
