import sys
import matplotlib.pyplot as plt
import cv2
import os
import numpy as np
from pyntcloud import PyntCloud
import subprocess


def find_global_min_max(pcd_folder_path):
    min_x, max_x = np.inf, -np.inf
    min_y, max_y = np.inf, -np.inf
    min_z, max_z = np.inf, -np.inf
    pcd_files = [os.path.join(pcd_folder_path, f) for f in os.listdir(pcd_folder_path) if f.endswith('.pcd')]
    for pcd_path in pcd_files:
        cloud = PyntCloud.from_file(pcd_path)
        points = cloud.points
        min_x, max_x = min(min_x, points['x'].min()), max(max_x, points['x'].max())
        min_y, max_y = min(min_y, points['y'].min()), max(max_y, points['y'].max())
        min_z, max_z = min(min_z, points['z'].min()), max(max_z, points['z'].max())
    return (min_x, max_x), (min_y, max_y), (min_z, max_z)


def process_pcd_to_top_view_image(pcd_path, output_path, axis_limits):
    cloud = PyntCloud.from_file(pcd_path)
    points = cloud.points
    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(points['x'], points['y'], points['z'], s=1, color='white')
    ax.view_init(elev=90, azim=90)
    ax.set_xlim(axis_limits[0])
    ax.set_ylim(axis_limits[1])
    ax.set_zlim(axis_limits[2])
    ax.set_facecolor('black')
    fig.patch.set_facecolor('black')
    plt.axis('off')
    plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
    plt.close()


def create_video_from_images(image_folder, output_video_path, frame_rate):
    images = [img for img in sorted(os.listdir(image_folder)) if img.endswith(".jpg")]
    frame = cv2.imread(os.path.join(image_folder, images[0]))
    height, width, layers = frame.shape
    video = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), float(frame_rate), (width, height))
    for image in images:
        frame = cv2.imread(os.path.join(image_folder, image))
        video.write(frame)
    video.release()
    print(f"Video successfully saved to: {output_video_path}")
    convert_video(output_video_path)


def convert_video(input_video_path):
    output_video_path = input_video_path.replace(".mp4", "_converted.mp4")
    command = f"ffmpeg -i {input_video_path} -c:v libx264 -preset fast -crf 23 {output_video_path}"
    subprocess.run(command, shell=True)
    os.remove(input_video_path)
    print(f"Converted video saved as {output_video_path} and original video deleted.")


if __name__ == '__main__':
    bag_base_name = sys.argv[1]
    output_base_path = sys.argv[2]

    lidar_output_base = os.path.join(output_base_path, bag_base_name + '_pcd_lidar')
    pcd_folder_path = os.path.join(lidar_output_base, 'pcd_ascii')
    images_folder_path = os.path.join(lidar_output_base, 'images')
    video_path = os.path.join(lidar_output_base, 'output_video.mp4')
    frame_rate_file = os.path.join(output_base_path, bag_base_name + '_pcd_lidar', 'frame_rate.txt')

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

    with open(frame_rate_file, 'r') as file:
        frame_rate = file.read().strip()

    axis_limits = find_global_min_max(pcd_folder_path)

    pcd_files = [f for f in os.listdir(pcd_folder_path) if f.endswith('.pcd')]
    for file_name in pcd_files:
        pcd_file_path = os.path.join(pcd_folder_path, file_name)
        image_file_path = os.path.join(images_folder_path, file_name.replace('.pcd', '.jpg'))
        process_pcd_to_top_view_image(pcd_file_path, image_file_path, axis_limits)

    create_video_from_images(images_folder_path, video_path, frame_rate)
