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

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


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)


def process_file(file):
    a = view(0, 30, 0, 10, 0, 1)
    colormap = np.array([[128, 130, 120],
                         [235, 0, 205],
                         [0, 215, 0],
                         [235, 155, 0]]) / 255.0
    root = file[0]
    inputFilePath = root + "/" + file[1]
    tempFilePath = root[:-3] + "pcd_ascii/" + file[1]
    outputFilePath = root[:-3] + "jpg/" + file[1][:-3] + "jpg"

    command = "pcl_convert_pcd_ascii_binary " + inputFilePath + " " + tempFilePath + " 0"
    os.system(command)

    if not os.path.exists(tempFilePath):
        return

    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 not math.isnan(point_x):
            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)


def parse(input_dir, output_dir):
    s_time = datetime.datetime.now()
    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
        start_time = bag.get_start_time()
        end_time = bag.get_end_time()
        print('start_time=', start_time)
        print('end_time=', end_time)
        for topic, msg, t in bag.read_messages(topics=["/points_concat"]):
            num_count += 1
            # print(num_count)
        duration = end_time - start_time
    bagtime = duration
    hz = str(float(float(num_count) / bagtime))
    print(hz)

    #######################解析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)

    # Create a list of files to be processed
    files_to_process = []
    for root, dirs, files in os.walk(output_dir[:-4]):
        for file in files:
            if file.endswith(".pcd"):
                files_to_process.append((root, file))

    num_processes = 4

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

    ####################### 将转化的点云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()
    e_time = datetime.datetime.now()
    print("using " + str(e_time - s_time) + "s")
    shutil.rmtree(file1)
    shutil.rmtree(file2)
    shutil.rmtree(output_dir)
    return output_dir[0:-4], num_count
